annet 0.11__py3-none-any.whl → 0.12__py3-none-any.whl

annet/generators/base.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+import abc
+import contextlib
+import textwrap
+from typing import Union, List
+
+from annet import tabparser, tracing
+from annet.tracing import tracing_connector
+from .exceptions import InvalidValueFromGenerator
+
+
+class DefaultBlockIfCondition:
+    pass
+
+
+ParamsList = tabparser.JuniperList
+
+
+class GenStringable(abc.ABC):
+    @abc.abstractmethod
+    def gen_str(self) -> str:
+        pass
+
+
+def _filter_str(value: Union[
+    str, int, float, tabparser.JuniperList, ParamsList, GenStringable]):
+    if isinstance(value, (
+            str,
+            int,
+            float,
+            tabparser.JuniperList,
+            ParamsList,
+    )):
+        return str(value)
+
+    if hasattr(value, "gen_str") and callable(value.gen_str):
+        return value.gen_str()
+
+    raise InvalidValueFromGenerator(
+        "Invalid yield type: %s(%s)" % (type(value).__name__, value))
+
+
+def _split_and_strip(text):
+    if "\n" in text:
+        rows = textwrap.dedent(text).strip().split("\n")
+    else:
+        rows = [text]
+    return rows
+
+
+# =====
+class BaseGenerator:
+    TYPE: str
+    TAGS: List[str]
+
+    def supports_device(self, device) -> bool:  # pylint: disable=unused-argument
+        return True
+
+
+class TreeGenerator(BaseGenerator):
+    def __init__(self, indent="  "):
+        self._indents = []
+        self._rows = []
+        self._block_path = []
+        self._indent = indent
+
+    @tracing.contextmanager(min_duration="0.1")
+    @contextlib.contextmanager
+    def block(self, *tokens, indent=None):
+        span = tracing_connector.get().get_current_span()
+        if span:
+            span.set_attribute("tokens", " ".join(map(str, tokens)))
+
+        indent = self._indent if indent is None else indent
+        block = " ".join(map(_filter_str, tokens))
+        self._block_path.append(block)
+        self._append_text(block)
+        self._indents.append(indent)
+        yield
+        self._indents.pop(-1)
+        self._block_path.pop(-1)
+
+    @contextlib.contextmanager
+    def block_if(self, *tokens, condition=DefaultBlockIfCondition):
+        if condition is DefaultBlockIfCondition:
+            condition = (None not in tokens and "" not in tokens)
+        if condition:
+            with self.block(*tokens):
+                yield
+                return
+        yield
+
+    @contextlib.contextmanager
+    def multiblock(self, *blocks):
+        if blocks:
+            blk = blocks[0]
+            tokens = blk if isinstance(blk, (list, tuple)) else [blk]
+            with self.block(*tokens):
+                with self.multiblock(*blocks[1:]):
+                    yield
+                    return
+        yield
+
+    @contextlib.contextmanager
+    def multiblock_if(self, *blocks, condition=DefaultBlockIfCondition):
+        if condition is DefaultBlockIfCondition:
+            condition = (None not in blocks)
+            if condition:
+                if blocks:
+                    blk = blocks[0]
+                    tokens = blk if isinstance(blk, (list, tuple)) else [blk]
+                    with self.block(*tokens):
+                        with self.multiblock(*blocks[1:]):
+                            yield
+                            return
+        yield
+
+    # ===
+    def _append_text(self, text):
+        self._append_text_cb(text)
+
+    def _append_text_cb(self, text, row_cb=None):
+        for row in _split_and_strip(text):
+            if row_cb:
+                row = row_cb(row)
+            self._rows.append("".join(self._indents) + row)
+
+
+class TextGenerator(TreeGenerator):
+    def __add__(self, line):
+        self._append_text(line)
+        return self
+
+    def __iter__(self):
+        yield from self._rows

annet/generators/common/initial.py

@@ -14,7 +14,7 @@ class InitialConfig(PartialGenerator):
     """
     def __init__(self, storage=None):
         self._do_run = not storage
-        super().__init__(storage=None)
+        super().__init__(storage=storage)
 
     def run_huawei(self, device):
         if not self._do_run:

annet/generators/entire.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import pkgutil
+import re
+import types
+from typing import (
+    FrozenSet,
+    Iterable,
+    List,
+    Optional,
+    Set,
+    Union,
+)
+
+from annet.lib import (
+    flatten,
+    jinja_render,
+    mako_render,
+)
+from .base import BaseGenerator, _filter_str
+from .exceptions import NotSupportedDevice
+
+
+class Entire(BaseGenerator):
+    TYPE = "ENTIRE"
+    TAGS: List[str] = []
+    REQUIRED_PACKAGES: FrozenSet[str] = frozenset()
+
+    def __init__(self, storage):
+        self.storage = storage
+        # между генераторами для одного и того же path - выбирается тот что больше
+        if not hasattr(self, "prio"):
+            self.prio = 100
+        self.__device = None
+
+    def supports_device(self, device):
+        return bool(self.path(device))
+
+    def run(self, device) -> Union[None, str, Iterable[Union[str, tuple]]]:
+        raise NotImplementedError
+
+    def reload(self, device) -> Optional[
+        str]:  # pylint: disable=unused-argument
+        return
+
+    def get_reload_cmds(self, device) -> str:
+        ret = self.reload(device) or ""
+        path = self.path(device)
+        if path and device.hw.PC and device.hw.soft.startswith(
+            ("Cumulus", "SwitchDev", "SONiC"),
+        ):
+            parts = []
+            if ret:
+                parts.append(ret)
+            parts.append("/usr/bin/etckeeper commitreload %s" % path)
+            return "\n".join(parts)
+        return ret
+
+    def path(self, device) -> Optional[str]:
+        raise NotImplementedError("Required PATH for ENTIRE generator")
+
+    # pylint: disable=unused-argument
+    def is_safe(self, device) -> bool:
+        """Output gen results when --acl-safe flag is used"""
+        return False
+
+    def read(self, path) -> str:
+        return pkgutil.get_data(__name__, path).decode()
+
+    def mako(self, text, **kwargs) -> str:
+        return mako_render(text, dedent=True, device=self.__device, **kwargs)
+
+    def jinja(self, text, **kwargs) -> str:
+        return jinja_render(text, dedent=True, device=self.__device, **kwargs)
+
+    # =====
+
+    @classmethod
+    def get_aliases(cls) -> Set[str]:
+        return {cls.__name__, *cls.TAGS}
+
+    def __call__(self, device):
+        self.__device = device
+        parts = []
+        run_res = self.run(device)
+        if isinstance(run_res, str):
+            run_res = (run_res,)
+        if run_res is None or not isinstance(run_res, (tuple, types.GeneratorType)):
+            raise Exception("generator %s returns %s" % (
+                self.__class__.__name__, type(run_res)))
+        for text in run_res:
+            if isinstance(text, tuple):
+                text = " ".join(map(_filter_str, flatten(text)))
+            assert re.search(r"\bNone\b", text) is None, \
+                "Found 'None' in yield result: %s" % text
+            parts.append(text)
+        return "\n".join(parts)

annet/generators/exceptions.py

@@ -0,0 +1,10 @@
+class GeneratorError(Exception):
+    pass
+
+
+class InvalidValueFromGenerator(ValueError):
+    pass
+
+
+class NotSupportedDevice(GeneratorError):
+    pass

annet/generators/jsonfragment.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+import contextlib
+from typing import (
+    Any,
+    Dict,
+    List,
+    Optional,
+    Set,
+    Union,
+)
+
+from annet.storage import Device, Storage
+from .base import TreeGenerator, _filter_str
+from .exceptions import NotSupportedDevice
+
+
+class JSONFragment(TreeGenerator):
+    """Generates parts of JSON config file."""
+
+    TYPE = "JSON_FRAGMENT"
+    TAGS: List[str] = []
+
+    def __init__(self, storage: Storage):
+        super().__init__()
+        self.storage = storage
+        self._json_config: Dict[str, Any] = {}
+        self._config_pointer: List[str] = []
+
+        # if two generators edit same file, commands from generator with greater `reload_prio` will be used
+        if not hasattr(self, "reload_prio"):
+            self.reload_prio = 100
+
+    def supports_device(self, device: Device):
+        return bool(self.path(device))
+
+    def path(self, device: Device) -> Optional[str]:
+        raise NotImplementedError("Required PATH for JSON_FRAGMENT generator")
+
+    @classmethod
+    def get_aliases(cls) -> Set[str]:
+        return {cls.__name__, *cls.TAGS}
+
+    def acl(self, device: Device) -> Union[str, List[str]]:
+        """
+        Restrict the generator to a specified ACL using JSON Pointer syntax.
+
+        Expected ACL to be a list of strings, but a single string is also allowed.
+        """
+        raise NotImplementedError("Required ACL for JSON_FRAGMENT generator")
+
+    def acl_safe(self, device: Device) -> Union[str, List[str]]:
+        """
+        Restrict the generator to a specified safe ACL using JSON Pointer syntax.
+
+        Expected ACL to be a list of strings, but a single string is also allowed.
+        """
+        raise NotImplementedError("Required ACL for JSON_FRAGMENT generator")
+
+    def run(self, device: Device):
+        raise NotImplementedError
+
+    def get_reload_cmds(self, device: Device) -> str:
+        ret = self.reload(device) or ""
+        return ret
+
+    def reload(self, device) -> Optional[str]:
+        raise NotImplementedError
+
+    @contextlib.contextmanager
+    def block(self, *tokens, indent=None):  # pylint: disable=unused-argument
+        block_str = "".join(map(_filter_str, tokens))
+        self._config_pointer.append(block_str)
+        try:
+            yield
+        finally:
+            self._config_pointer.pop()
+
+    @contextlib.contextmanager
+    def block_piped(self, *tokens, indent=None):  # pylint: disable=unused-argument
+        block_str = "|".join(map(_filter_str, tokens))
+        self._config_pointer.append(block_str)
+        try:
+            yield
+        finally:
+            self._config_pointer.pop()
+
+    def __call__(self, device: Device, annotate: bool = False):
+        for cfg_fragment in self.run(device):
+            self._set_or_replace_dict(self._config_pointer, cfg_fragment)
+        return self._json_config
+
+    def _set_or_replace_dict(self, pointer, value):
+        if not pointer:
+            if self._json_config == {}:
+                self._json_config = value
+            else:
+                self._json_config = [self._json_config, value]
+        else:
+            self._set_dict(self._json_config, pointer, value)
+
+    @classmethod
+    def _to_str(cls, value: Any) -> str:
+        if isinstance(value, str):
+            return value
+        elif isinstance(value, list):
+            return [cls._to_str(x) for x in value]
+        elif isinstance(value, dict):
+            for k, v in value.items():
+                value[k] = cls._to_str(v)
+            return value
+        return str(value)
+
+    @classmethod
+    def _set_dict(cls, cfg, pointer, value):
+        # pointer has at least one key
+        if len(pointer) == 1:
+            if pointer[0] in cfg:
+                cfg[pointer[0]] = [cfg[pointer[0]], cls._to_str(value)]
+            else:
+                cfg[pointer[0]] = cls._to_str(value)
+        else:
+            if pointer[0] not in cfg:
+                cfg[pointer[0]] = {}
+            cls._set_dict(cfg[pointer[0]], pointer[1:], cls._to_str(value))

annet/generators/partial.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+import re
+from typing import (
+    Iterable,
+    List,
+    Set,
+    Union,
+)
+
+from annet.lib import (
+    add_annotation,
+    flatten,
+)
+from .base import TreeGenerator, _filter_str
+from .exceptions import NotSupportedDevice
+
+
+class PartialGenerator(TreeGenerator):
+    TYPE = "PARTIAL"
+    TAGS: List[str] = []
+
+    def __init__(self, storage):
+        super().__init__()
+        self.storage = storage
+        self._annotate = False
+        self._running_gen = None
+        self._annotations = []
+        self._annotation_module = self.__class__.__module__ or ""
+
+    def supports_device(self, device) -> bool:
+        if self.__class__.run is PartialGenerator.run:
+            return bool(self._get_vendor_func(device.hw.vendor, "run"))
+        else:
+            return True
+
+    def acl(self, device):
+        acl_func = self._get_vendor_func(device.hw.vendor, "acl")
+        if acl_func:
+            return acl_func(device)
+
+    def acl_safe(self, device):
+        acl_func = self._get_vendor_func(device.hw.vendor, "acl_safe")
+        if acl_func:
+            return acl_func(device)
+
+    def run(self, device) -> Iterable[Union[str, tuple]]:
+        run_func = self._get_vendor_func(device.hw.vendor, "run")
+        if run_func:
+            return run_func(device)
+
+    def get_user_runner(self, device):
+        if self.__class__.run is not PartialGenerator.run:
+            return self.run
+        return self._get_vendor_func(device.hw.vendor, "run")
+
+    def _get_vendor_func(self, vendor: str, func_name: str):
+        attr_name = f"{func_name}_{vendor}"
+        if hasattr(self, attr_name):
+            return getattr(self, attr_name)
+        return None
+
+    # =====
+
+    @classmethod
+    def get_aliases(cls) -> Set[str]:
+        return {cls.__name__, *cls.TAGS}
+
+    def __call__(self, device, annotate=False):
+        self._indents = []
+        self._rows = []
+        self._running_gen = self.run(device)
+        self._annotate = annotate
+
+        if annotate and self.__class__.__module__:
+            self._annotation_module = ".".join(
+                self.__class__.__module__.split(".")[-2:])
+
+        for text in self._running_gen:
+            if isinstance(text, tuple):
+                text = " ".join(map(_filter_str, flatten(text)))
+            else:
+                text = _filter_str(text)
+            self._append_text(text)
+
+        for row in self._rows:
+            assert re.search(r"\bNone\b", row) is None, \
+                "Found 'None' in yield result: %s" % (row)
+        if annotate:
+            generated_rows = (
+                add_annotation(x, y)
+                for (x, y) in zip(self._rows, self._annotations)
+            )
+        else:
+            generated_rows = self._rows
+        return "\n".join(generated_rows) + "\n"
+
+    def _append_text(self, text):
+        def annotation_cb(row):
+            annotation = "%s:%d" % self.get_running_line()
+            self._annotations.append(annotation)
+            return row
+
+        self._append_text_cb(
+            text,
+            annotation_cb if self._annotate else None
+        )
+
+    def get_running_line(self):
+        if not self._running_gen or not self._running_gen.gi_frame:
+            return (repr(self._running_gen), -1)
+        return self._annotation_module, self._running_gen.gi_frame.f_lineno
+
+    @classmethod
+    def literal(cls, item):
+        return '"{}"'.format(item)
+
+    def __repr__(self):
+        return "<%s>" % self.__class__.__name__

annet/generators/perf.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import time
+from typing import (
+    Optional,
+    Union,
+)
+
+from annet import tracing
+from annet.tracing import tracing_connector
+from annet.types import (
+    GeneratorPartialRunArgs,
+    GeneratorPerf,
+)
+from .entire import Entire
+from .partial import PartialGenerator
+
+
+class GeneratorPerfMesurer:
+    def __init__(
+            self,
+            gen: Union[PartialGenerator, Entire],
+            run_args: Optional[GeneratorPartialRunArgs] = None,
+            trace_min_duration: tracing.MinDurationT = None
+    ):
+        self._gen = gen
+        self._run_args = run_args
+
+        self._start_time: float = 0.0
+        self._span_ctx = None
+        self._span = None
+        self._trace_min_duration = trace_min_duration
+
+        self.last_result: Optional[GeneratorPerf] = None
+
+    def start(self) -> None:
+        self.last_result = None
+
+        self._span_ctx = tracing_connector.get().start_as_current_span(
+            "gen:call",
+            tracer_name=self._gen.__class__.__module__,
+            min_duration=self._trace_min_duration,
+        )
+        self._span = self._span_ctx.__enter__()  # pylint: disable=unnecessary-dunder-call
+
+        if self._span:
+            self._span.set_attributes(
+                {"generator.class": self._gen.__class__.__name__})
+            if self._run_args:
+                tracing_connector.get().set_device_attributes(
+                    self._span, self._run_args.device,
+                )
+
+        self._start_time = time.monotonic()
+
+    def finish(self, exc_type=None, exc_val=None,
+               exc_tb=None) -> GeneratorPerf:
+        total = time.monotonic() - self._start_time
+        self._span_ctx.__exit__(exc_type, exc_val, exc_tb)
+
+        meta = {}
+        if tracing_connector.get().enabled:
+            span_context = self._span.get_span_context()
+            meta = {
+                "span": {
+                    "trace_id": str(span_context.trace_id),
+                    "span_id": str(span_context.span_id),
+                },
+            }
+
+        self.last_result = GeneratorPerf(total=total, rt=total, meta=meta)
+        return self.last_result
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.finish(exc_type, exc_val, exc_tb)

annet/generators/ref.py

@@ -0,0 +1,15 @@
+from .partial import PartialGenerator
+
+
+class RefGenerator(PartialGenerator):
+    def __init__(self, storage, groups=None):
+        super().__init__(storage)
+        self.groups = groups
+
+    def ref(self, device):
+        if hasattr(self, "ref_" + device.hw.vendor):
+            return getattr(self, "ref_" + device.hw.vendor)(device)
+        return ""
+
+    def with_groups(self, groups):
+        return type(self)(self.storage, groups)