1
"""
2
Utilities for CLI programs
3
"""
4

5 4
import argparse
6 4
import copy
7 4
import importlib
8 4
import json
9 4
import signal
10 4
from functools import partial
11

12 4
import yaml
13

14

15 4
def import_module(module, package=None):
16
    """Protected import of a module"""
17 1
    try:
18 1
        ret = importlib.import_module(module, package=package)
19 0
    except ModuleNotFoundError:
20 0
        if package is not None:
21 0
            raise ModuleNotFoundError("Requested module/package '{}/{}' not found.".format(module, package))
22 0
        raise ModuleNotFoundError("Requested module '{}' not found.".format(module))
23 1
    return ret
24

25

26 4
def read_config_file(fname):
27
    """Reads a JSON or YAML file."""
28 4
    if fname.endswith(".yaml") or fname.endswith(".yml"):
29 4
        try:
30 4
            rfunc = partial(yaml.load, Loader=yaml.FullLoader)
31 0
        except AttributeError:
32 0
            rfunc = yaml.load
33 0
    elif fname.endswith(".json"):
34 0
        rfunc = json.load
35
    else:
36 0
        raise TypeError("Did not understand file type {}.".format(fname))
37

38 4
    try:
39 4
        with open(fname, "r") as handle:
40 4
            ret = rfunc(handle)
41 0
    except FileNotFoundError:
42 0
        raise FileNotFoundError("No config file found at {}.".format(fname))
43

44 4
    return ret
45

46

47 4
def argparse_config_merge(parser, parsed_options, config_options, parser_default=None, check=True):
48
    """Merges options between a configuration file and a parser
49

50
    Parameters
51
    ----------
52
    parser : ArgumentParser
53
    config_options : dict
54
    parser_default : List, Optional
55
    """
56 4
    config_options = copy.deepcopy(config_options)
57

58 4
    if check:
59 0
        default_options = vars(parser.parse_args(args=parser_default))
60 0
        diff = config_options.keys() - default_options.keys()
61 0
        if diff:
62 0
            raise argparse.ArgumentError(
63
                None, "Unknown arguments found in configuration file: {}.".format(", ".join(diff))
64
            )
65

66
    # Add in parsed options
67 4
    for k, v in parsed_options.items():
68

69
        # User has overridden default options, update config
70 4
        if v != parser.get_default(k):
71 4
            config_options[k] = v
72

73
        # Add in missing defaults
74 4
        if k not in config_options:
75 4
            config_options[k] = v
76

77 4
    return config_options
78

79

80 4
def install_signal_handlers(loop, cleanup):
81
    """
82
    Install cleanup handlers to shutdown on:
83
     - Keyboard Interupt (SIGINT)
84
     - Shutdown kill/pkill (SIGTERM)
85
    """
86

87 4
    old_handlers = {}
88

89 4
    def handle_signal(sig, frame):
90 4
        async def cleanup_and_stop():
91 4
            try:
92 4
                cleanup()
93
            finally:
94 4
                loop.stop()
95

96 4
        loop.add_callback_from_signal(cleanup_and_stop)
97

98
        # Add old handlers back in so we do not cleanup twice
99 4
        signal.signal(sig, old_handlers[sig])
100

101 4
    for sig in [signal.SIGINT, signal.SIGTERM]:
102 4
        old_handlers[sig] = signal.signal(sig, handle_signal)

Read our documentation on viewing source code .

Loading