tools.py 4.17 KB

# python standard library
import asyncio
import logging
from os import path
import subprocess
from typing import Any, List

# third party libraries
import yaml


# setup logger for this module
logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# load data from yaml file
# ---------------------------------------------------------------------------
def load_yaml(filename: str, default: Any = None) -> Any:
    filename = path.expanduser(filename)
    try:
        f = open(filename, 'r', encoding='utf-8')
    except FileNotFoundError:
        logger.error(f'Cannot open "{filename}": not found')
    except PermissionError:
        logger.error(f'Cannot open "{filename}": no permission')
    except OSError:
        logger.error(f'Cannot open file "{filename}"')
    else:
        with f:
            try:
                default = yaml.safe_load(f)
            except yaml.YAMLError as e:
                if hasattr(e, 'problem_mark'):
                    mark = e.problem_mark
                    logger.error(f'File "{filename}" near line {mark.line}, '
                                 f'column {mark.column+1}')
                else:
                    logger.error(f'File "{filename}"')
    finally:
        return default


# ---------------------------------------------------------------------------
# Runs a script and returns its stdout parsed as yaml, or None on error.
# The script is run in another process but this function blocks waiting
# for its termination.
# ---------------------------------------------------------------------------
def run_script(script: str,
               args: List[str] = [],
               stdin: str = '',
               timeout: int = 2) -> Any:

    script = path.expanduser(script)
    try:
        cmd = [script] + [str(a) for a in args]
        p = subprocess.run(cmd,
                           input=stdin,
                           stdout=subprocess.PIPE,
                           stderr=subprocess.STDOUT,
                           universal_newlines=True,
                           timeout=timeout,
                           )
    except FileNotFoundError:
        logger.error(f'Can not execute script "{script}": not found.')
    except PermissionError:
        logger.error(f'Can not execute script "{script}": wrong permissions.')
    except OSError:
        logger.error(f'Can not execute script "{script}": unknown reason.')
    except subprocess.TimeoutExpired:
        logger.error(f'Timeout {timeout}s exceeded while running "{script}".')
    except Exception:
        logger.error(f'An Exception ocurred running {script}.')
    else:
        if p.returncode != 0:
            logger.error(f'Return code {p.returncode} running "{script}".')
        else:
            try:
                output = yaml.safe_load(p.stdout)
            except Exception:
                logger.error(f'Error parsing yaml output of "{script}"')
            else:
                return output


# ----------------------------------------------------------------------------
# Same as above, but asynchronous
# ----------------------------------------------------------------------------
async def run_script_async(script: str,
                           args: List[str] = [],
                           stdin: str = '',
                           timeout: int = 2) -> Any:

    script = path.expanduser(script)
    args = [str(a) for a in args]

    p = await asyncio.create_subprocess_exec(
        script, *args,
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.DEVNULL,
        )

    try:
        stdout, stderr = await asyncio.wait_for(
            p.communicate(input=stdin.encode('utf-8')),
            timeout=timeout
            )
    except asyncio.TimeoutError:
        logger.warning(f'Timeout {timeout}s running script "{script}".')
        return

    if p.returncode != 0:
        logger.error(f'Return code {p.returncode} running "{script}".')
    else:
        try:
            output = yaml.safe_load(stdout.decode('utf-8', 'ignore'))
        except Exception:
            logger.error(f'Error parsing yaml output of "{script}"')
        else:
            return output