<commit_before>import os

from loguru import logger

from flexget import plugin
from flexget.event import event

logger = logger.bind(name='free_space')


def get_free_space(folder):
    """ Return folder/drive free space (in megabytes)"""
    if os.name == 'nt':
        import ctypes

        free_bytes = ctypes.c_ulonglong(0)
        ctypes.windll.kernel32.GetDiskFreeSpaceExW(
            ctypes.c_wchar_p(folder), None, None, ctypes.pointer(free_bytes)
        )
        return free_bytes.value / (1024 * 1024)
    else:
        stats = os.statvfs(folder)
        return (stats.f_bavail * stats.f_frsize) / (1024 * 1024)


class PluginFreeSpace:
    """Aborts a task if an entry is accepted and there is less than a certain amount of space free on a drive."""

    schema = {
        'oneOf': [
            {'type': 'number'},
            {
                'type': 'object',
                'properties': {
                    'space': {'type': 'number'},
                    'path': {'type': 'string', 'format': 'path'},
                },
                'required': ['space'],
                'additionalProperties': False,
            },
        ]
    }

    def prepare_config(self, config, task):
        if isinstance(config, (float, int)):
            config = {'space': config}
        # Use config path if none is specified
        if not config.get('path'):
            config['path'] = task.manager.config_base
        return config

    @plugin.priority(plugin.PRIORITY_FIRST)
    def on_task_download(self, task, config):
        config = self.prepare_config(config, task)
        # Only bother aborting if there were accepted entries this run.
        if task.accepted:
            if get_free_space(config['path']) < config['space']:
                logger.error(
                    'Less than {} MB of free space in {} aborting task.',
                    config['space'],
                    config['path'],
                )
                # backlog plugin will save and restore the task content, if available<commit_msg>Add ability to poll remote drives.<commit_after>import os

from loguru import logger

from flexget import plugin
from flexget.event import event

logger = logger.bind(name='free_space')


def get_free_space(config):
    """ Return folder/drive free space (in megabytes)"""
    if config['remote']:
        import paramiko
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            ssh.connect(config['host'], config['port'], config['user'], key_filename=config['ssh_key_filepath'], look_for_keys=True, timeout=5000)
        except Exception as e:
            logger.error("Issue connecting to remote host. {}", e)
            return "ERR"
        if config['allotment'] != -1:
            stdin, stdout, stderr = ssh.exec_command(f"du -s {config['path']} | cut -f 1")
        else:
            stdin, stdout, stderr = ssh.exec_command(f"df -k {config['path']} | tail -1 | tr -s ' ' | cut -d' ' -f4")
        outlines = stdout.readlines()
        resp = ''.join(outlines)
        try:
            if config['allotment'] != -1:
                free = int(config['allotment']) - ((int(resp.strip()) * 1024) / 1000000)
            else:
                free = (int(resp.strip()) / 1000)
        except ValueError:
            logger.error('Non-integer was returned when calcualting disk usage.')
            task.abort('Remote disk usage could not be calculated.')
        ssh.close()
        return free
    elif os.name == 'nt':
        import ctypes
        free_bytes = ctypes.c_ulonglong(0)
        ctypes.windll.kernel32.GetDiskFreeSpaceExW(
            ctypes.c_wchar_p(config['path']), None, None, ctypes.pointer(free_bytes)
        )
        return free_bytes.value / (1024 * 1024)
    else:
        stats = os.statvfs(config['path'])
        return (stats.f_bavail * stats.f_frsize) / (1024 * 1024)


class PluginFreeSpace:
    """Aborts a task if an entry is accepted and there is less than a certain amount of space free on a drive."""

    schema = {
        'oneOf': [
            {'type': 'number'},
            {
                'type': 'object',
                'properties': {
                    'space': {'type': 'number'},
                    'path': {'type': 'string'},
                    'remote': {'type': 'boolean', 'default': False},
                    'port': {'type': 'integer', 'default': 22},
                    'host': {'type': 'string'},
                    'user': {'type': 'string'},
                    'ssh_key_filepath': {'type': 'string'},
                    'allotment': {'type': 'number', 'default': -1},
                },
                'required': ['space'],
                'dependencies': {
                    'remote': ['host', 'user', 'ssh_key_filepath', 'path']
                },
                'additionalProperties': False,
            },
        ]
    }

    def prepare_config(self, config, task):
        if isinstance(config, (float, int)):
            config = {'space': config}
        # Use config path if none is specified
        if not config.get('path'):
            config['path'] = task.manager.config_base
        return config

    @plugin.priority(plugin.PRIORITY_FIRST)
    def on_task_download(self, task, config):
        config = self.prepare_config(config, task)
        # Only bother aborting if there were accepted entries this run.
        if task.accepted:
            free_space = get_free_space(config)
            if free_space == "ERR":
                task.abort('Could not connect to remote host.')
            elif free_space < config['space']:
                logger.error(
                    'Less than {} MB of free space in {} aborting task.',
                    config['space'],
                    config['path'],
                )
                # backlog plugin will save and restore the task content, if available
{'fraction_non_alphanumeric': 0.12588592642591967, 'fraction_numerical': 0.010968612892338846, 'mean_word_length': 1.8757884522076662, 'pattern_counts': {'":': 1, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4807382', 'n_tokens_mistral': 1716, 'n_tokens_neox': 1609, 'n_words': 455}