File size: 7,090 Bytes

fbd9366

import codecs
from collections import Counter
import fnmatch
import hashlib
import os
import pickle
from typing import Any, Callable, Dict, List, Optional, Union

from typing_extensions import Literal


def set_os_envs(envs: Optional[Dict[str, Any]] = None):
    """
    Special value __delete__ or None indicates that the ENV_VAR should be removed
    """
    if envs is None:
        envs = {}
    DEL = {None, "__delete__"}
    for k, v in envs.items():
        if v in DEL:
            os.environ.pop(k, None)
    os.environ.update({k: str(v) for k, v in envs.items() if v not in DEL})


def argmax(L):
    return max(zip(L, range(len(L))))[1]


def _match_patterns_helper(element, patterns):
    for p in patterns:
        if callable(p) and p(element):
            return True
        if fnmatch.fnmatch(element, p):
            return True
    return False


def match_patterns(
    item: str,
    include: Union[str, List[str], Callable, List[Callable], None] = None,
    exclude: Union[str, List[str], Callable, List[Callable], None] = None,
    *,
    precedence: Literal["include", "exclude"] = "exclude",
):
    """
    Args:
        include: None to disable `include` filter and delegate to exclude
        precedence: "include" or "exclude"
    """
    assert precedence in ["include", "exclude"]
    if exclude is None:
        exclude = []
    if isinstance(exclude, (str, Callable)):
        exclude = [exclude]
    if isinstance(include, (str, Callable)):
        include = [include]
    if include is None:
        # exclude is the sole veto vote
        return not _match_patterns_helper(item, exclude)

    if precedence == "include":
        return _match_patterns_helper(item, include)
    else:
        if _match_patterns_helper(item, exclude):
            return False
        else:
            return _match_patterns_helper(item, include)


def filter_patterns(
    items: List[str],
    include: Union[str, List[str], Callable, List[Callable], None] = None,
    exclude: Union[str, List[str], Callable, List[Callable], None] = None,
    *,
    precedence: Literal["include", "exclude"] = "exclude",
    ordering: Literal["original", "include"] = "original",
):
    """
    Args:
        ordering: affects the order of items in the returned list. Does not affect the
            content of the returned list.
            - "original": keep the ordering of items in the input list
            - "include": order items by the order of include patterns
    """
    assert ordering in ["original", "include"]
    if include is None or isinstance(include, str) or ordering == "original":
        return [
            x
            for x in items
            if match_patterns(x, include=include, exclude=exclude, precedence=precedence)
        ]
    else:
        items = items.copy()
        ret = []
        for inc in include:
            for i, item in enumerate(items):
                if item is None:
                    continue
                if match_patterns(item, include=inc, exclude=exclude, precedence=precedence):
                    ret.append(item)
                    items[i] = None
        return ret


def getitem_nested(cfg, key: str):
    """
    Recursively get key, if key has '.' in it
    """
    keys = key.split(".")
    for k in keys:
        assert k in cfg, f'{k} in key "{key}" does not exist in config'
        cfg = cfg[k]
    return cfg


def setitem_nested(cfg, key: str, value):
    """
    Recursively get key, if key has '.' in it
    """
    keys = key.split(".")
    for k in keys[:-1]:
        assert k in cfg, f'{k} in key "{key}" does not exist in config'
        cfg = cfg[k]
    cfg[keys[-1]] = value


def getattr_nested(obj, key: str):
    """
    Recursively get attribute
    """
    keys = key.split(".")
    for k in keys:
        assert hasattr(obj, k), f'{k} in attribute "{key}" does not exist'
        obj = getattr(obj, k)
    return obj


def setattr_nested(obj, key: str, value):
    """
    Recursively set attribute
    """
    keys = key.split(".")
    for k in keys[:-1]:
        assert hasattr(obj, k), f'{k} in attribute "{key}" does not exist'
        obj = getattr(obj, k)
    setattr(obj, keys[-1], value)


class PeriodicEvent:
    """
    triggers every period
    """

    def __init__(self, period: int, initial_value=0):
        self._period = period
        assert self._period >= 1
        self._last_threshold = initial_value
        self._last_value = initial_value
        self._trigger_counts = 0

    def __call__(self, new_value=None, increment=None):
        assert bool(new_value is None) != bool(increment is None), (
            "you must specify one and only one of new_value or increment, " "but not both"
        )
        d = self._period
        if new_value is None:
            new_value = self._last_value + increment
        assert new_value >= self._last_value, (
            f"value must be monotonically increasing. "
            f"Current value {new_value} < last value {self._last_value}"
        )
        self._last_value = new_value
        if new_value - self._last_threshold >= d:
            self._last_threshold += (new_value - self._last_threshold) // d * d
            self._trigger_counts += 1
            return True
        else:
            return False

    @property
    def trigger_counts(self):
        return self._trigger_counts

    @property
    def current_value(self):
        return self._last_value


class Once:
    def __init__(self):
        self._triggered = False

    def __call__(self):
        if not self._triggered:
            self._triggered = True
            return True
        else:
            return False

    def __bool__(self):
        raise RuntimeError("`Once` objects should be used by calling ()")


_GLOBAL_ONCE_SET = set()
_GLOBAL_NTIMES_COUNTER = Counter()


def global_once(name):
    """
    Try this to automate the name:
    https://gist.github.com/techtonik/2151727#gistcomment-2333747
    """
    if name in _GLOBAL_ONCE_SET:
        return False
    else:
        _GLOBAL_ONCE_SET.add(name)
        return True


def global_n_times(name, n: int):
    """
    Triggers N times
    """
    assert n >= 1
    if _GLOBAL_NTIMES_COUNTER[name] < n:
        _GLOBAL_NTIMES_COUNTER[name] += 1
        return True
    else:
        return False


class Every:
    def __init__(self, n: int, on_first: bool = False):
        assert n > 0
        self._i = 0 if on_first else 1
        self._n = n

    def __call__(self):
        return self._i % self._n == 0

    def __bool__(self):
        raise RuntimeError("`Every` objects should be used by calling ()")


def encode_base64(obj) -> str:
    return codecs.encode(pickle.dumps(obj), "base64").decode()


def decode_base64(s: str):
    return pickle.loads(codecs.decode(s.encode(), "base64"))


def safe_hash(input_tuple):
    # keep 128 bits of the hash
    tuple_string = repr(input_tuple).encode("utf-8")
    sha256 = hashlib.sha256()
    sha256.update(tuple_string)

    seed = int(sha256.hexdigest(), 16)

    return seed & 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF