diff --git a/nornir/core/configuration.py b/nornir/core/configuration.py index 8a04c02f..be5f12b9 100644 --- a/nornir/core/configuration.py +++ b/nornir/core/configuration.py @@ -111,6 +111,30 @@ def dict(self) -> Dict[str, Any]: } +class InventoryDataConfig(object): + __slots__ = "plugin", "options" + + class Parameters: + plugin = Parameter( + typ=str, default="InventoryDataDict", envvar="NORNIR_INVENTORY_DATA_PLUGIN" + ) + options = Parameter(default={}, envvar="NORNIR_INVENTORY_DATA_OPTIONS") + + def __init__( + self, + plugin: Optional[str] = None, + options: Optional[Dict[str, Any]] = None, + ) -> None: + self.plugin = self.Parameters.plugin.resolve(plugin) + self.options = self.Parameters.options.resolve(options) or {} + + def dict(self) -> Dict[str, Any]: + return { + "plugin": self.plugin, + "options": self.options, + } + + class LoggingConfig(object): __slots__ = "enabled", "level", "log_file", "format", "to_console", "loggers" @@ -245,6 +269,7 @@ class Config(object): "runner", "ssh", "inventory", + "inventory_data", "logging", "user_defined", ) @@ -252,6 +277,7 @@ class Config(object): def __init__( self, inventory: Optional[InventoryConfig] = None, + inventory_data: Optional[InventoryDataConfig] = None, ssh: Optional[SSHConfig] = None, logging: Optional[LoggingConfig] = None, core: Optional[CoreConfig] = None, @@ -259,6 +285,7 @@ def __init__( user_defined: Optional[Dict[str, Any]] = None, ) -> None: self.inventory = inventory or InventoryConfig() + self.inventory_data = inventory_data or InventoryDataConfig() self.ssh = ssh or SSHConfig() self.logging = logging or LoggingConfig() self.core = core or CoreConfig() @@ -269,6 +296,7 @@ def __init__( def from_dict( cls, inventory: Optional[Dict[str, Any]] = None, + inventory_data: Optional[Dict[str, Any]] = None, ssh: Optional[Dict[str, Any]] = None, logging: Optional[Dict[str, Any]] = None, core: Optional[Dict[str, Any]] = None, @@ -277,6 +305,7 @@ def from_dict( ) -> "Config": return cls( inventory=InventoryConfig(**inventory or {}), + inventory_data=InventoryDataConfig(**inventory_data or {}), ssh=SSHConfig(**ssh or {}), logging=LoggingConfig(**logging or {}), core=CoreConfig(**core or {}), @@ -289,6 +318,7 @@ def from_file( cls, config_file: str, inventory: Optional[Dict[str, Any]] = None, + inventory_data: Optional[Dict[str, Any]] = None, ssh: Optional[Dict[str, Any]] = None, logging: Optional[Dict[str, Any]] = None, core: Optional[Dict[str, Any]] = None, @@ -296,6 +326,7 @@ def from_file( user_defined: Optional[Dict[str, Any]] = None, ) -> "Config": inventory = inventory or {} + inventory_data = inventory_data or {} ssh = ssh or {} logging = logging or {} core = core or {} @@ -306,6 +337,9 @@ def from_file( data = yml.load(f) return cls( inventory=InventoryConfig(**{**data.get("inventory", {}), **inventory}), + inventory_data=InventoryDataConfig( + **{**data.get("inventory_data", {}), **inventory_data} + ), ssh=SSHConfig(**{**data.get("ssh", {}), **ssh}), logging=LoggingConfig(**{**data.get("logging", {}), **logging}), core=CoreConfig(**{**data.get("core", {}), **core}), @@ -316,6 +350,7 @@ def from_file( def dict(self) -> Dict[str, Any]: return { "inventory": self.inventory.dict(), + "inventory_data": self.inventory_data.dict(), "ssh": self.ssh.dict(), "logging": self.logging.dict(), "core": self.core.dict(), diff --git a/nornir/core/inventory.py b/nornir/core/inventory.py index f900dbff..bd51a5f0 100644 --- a/nornir/core/inventory.py +++ b/nornir/core/inventory.py @@ -16,10 +16,26 @@ from nornir.core.configuration import Config from nornir.core.exceptions import ConnectionAlreadyOpen, ConnectionNotOpen from nornir.core.plugins.connections import ConnectionPlugin, ConnectionPluginRegister +from nornir.core.plugins.inventory_data import ( + InventoryData, + InventoryDataPluginRegister, +) HostOrGroup = TypeVar("HostOrGroup", "Host", "Group") +def _init_inventory_data( + data: Optional[Dict[str, Any]] = None, configuration: Optional[Config] = None +) -> Union[Dict[str, Any], InventoryData]: + if not configuration: + configuration = Config() + InventoryDataPluginRegister.auto_register() + inventory_data_plugin = InventoryDataPluginRegister.get_plugin( + configuration.inventory_data.plugin + ) + return inventory_data_plugin(**configuration.inventory_data.options).load(data) + + class BaseAttributes(object): __slots__ = ("hostname", "port", "username", "password", "platform") @@ -125,9 +141,10 @@ def __init__( groups: Optional[ParentGroups] = None, data: Optional[Dict[str, Any]] = None, connection_options: Optional[Dict[str, ConnectionOptions]] = None, + configuration: Optional[Config] = None, ) -> None: self.groups = groups or ParentGroups() - self.data = data or {} + self.data = _init_inventory_data(data, configuration=configuration) self.connection_options = connection_options or {} super().__init__( hostname=hostname, @@ -208,8 +225,9 @@ def __init__( platform: Optional[str] = None, data: Optional[Dict[str, Any]] = None, connection_options: Optional[Dict[str, ConnectionOptions]] = None, + configuration: Optional[Config] = None, ) -> None: - self.data = data or {} + self.data = _init_inventory_data(data, configuration=configuration) self.connection_options = connection_options or {} super().__init__( hostname=hostname, @@ -252,6 +270,7 @@ def __init__( data: Optional[Dict[str, Any]] = None, connection_options: Optional[Dict[str, ConnectionOptions]] = None, defaults: Optional[Defaults] = None, + configuration: Optional[Config] = None, ) -> None: self.name = name self.defaults = defaults or Defaults(None, None, None, None, None, None, None) @@ -265,6 +284,7 @@ def __init__( groups=groups, data=data, connection_options=connection_options, + configuration=configuration, ) def extended_data(self) -> Dict[str, Any]: diff --git a/nornir/core/plugins/inventory_data.py b/nornir/core/plugins/inventory_data.py new file mode 100644 index 00000000..c182ede8 --- /dev/null +++ b/nornir/core/plugins/inventory_data.py @@ -0,0 +1,80 @@ +from typing import ( + Any, + Dict, + ItemsView, + KeysView, + Optional, + Protocol, + Type, + Union, + ValuesView, +) + +from nornir.core.plugins.register import PluginRegister + +INVENTORY_DATA_PLUGIN_PATH = "nornir.plugins.inventory_data" + + +class InventoryData(Protocol): + def __init__(self, *args: Any, **kwargs: Any) -> None: + """ + This method configures the plugin + """ + raise NotImplementedError("needs to be implemented by the plugin") + + def __getitem__(self, key: str) -> Any: + """ + This method configures the plugin + """ + raise NotImplementedError("needs to be implemented by the plugin") + + def get(self, key: str, default: Any = None) -> Any: + """ + This method configures the plugin + """ + raise NotImplementedError("needs to be implemented by the plugin") + + def __setitem__(self, key: str, value: Any) -> None: + """ + This method configures the plugin + """ + raise NotImplementedError("needs to be implemented by the plugin") + + def keys(self) -> KeysView[str]: + """ + This method configures the plugin + """ + raise NotImplementedError("needs to be implemented by the plugin") + + def values(self) -> ValuesView[Any]: + """ + This method configures the plugin + """ + raise NotImplementedError("needs to be implemented by the plugin") + + def items(self) -> ItemsView[str, Any]: + """ + This method configures the plugin + """ + raise NotImplementedError("needs to be implemented by the plugin") + + +class InventoryDataPlugin(Protocol): + def __init__(self, *args: Any, **kwargs: Any) -> None: + """ + This method configures the plugin + """ + raise NotImplementedError("needs to be implemented by the plugin") + + def load( + self, data: Optional[Dict[str, Any]] = None + ) -> Union[Dict[str, Any], InventoryData]: + """ + Returns the object containing the data + """ + raise NotImplementedError("needs to be implemented by the plugin") + + +InventoryDataPluginRegister: PluginRegister[Type[InventoryDataPlugin]] = PluginRegister( + INVENTORY_DATA_PLUGIN_PATH +) diff --git a/nornir/init_nornir.py b/nornir/init_nornir.py index 7de8f8a4..95a5d9e9 100644 --- a/nornir/init_nornir.py +++ b/nornir/init_nornir.py @@ -1,4 +1,5 @@ -from typing import Any +import inspect +from typing import Any, Union from nornir.core import Nornir from nornir.core.configuration import Config @@ -17,7 +18,18 @@ def load_inventory( ) -> Inventory: InventoryPluginRegister.auto_register() inventory_plugin = InventoryPluginRegister.get_plugin(config.inventory.plugin) - inv = inventory_plugin(**config.inventory.options).load() + inventory_plugin_params = config.inventory.options.copy() + + init_params = inspect.signature(inventory_plugin).parameters + if "configuration" in init_params: + config_parameter = init_params["configuration"] + if config_parameter.annotation is not inspect.Parameter.empty and ( + config_parameter.annotation == Config + or config_parameter.annotation == Union[Config, None] + ): + inventory_plugin_params.update({"configuration": config}) + + inv = inventory_plugin(**inventory_plugin_params).load() if config.inventory.transform_function: TransformFunctionRegister.auto_register() diff --git a/nornir/plugins/inventory/simple.py b/nornir/plugins/inventory/simple.py index 0ee0884c..4897f23c 100644 --- a/nornir/plugins/inventory/simple.py +++ b/nornir/plugins/inventory/simple.py @@ -1,9 +1,10 @@ import logging import pathlib -from typing import Any, Dict, Type +from typing import Any, Dict, Optional, Type import ruamel.yaml +from nornir.core.configuration import Config from nornir.core.inventory import ( ConnectionOptions, Defaults, @@ -33,7 +34,9 @@ def _get_connection_options(data: Dict[str, Any]) -> Dict[str, ConnectionOptions return cp -def _get_defaults(data: Dict[str, Any]) -> Defaults: +def _get_defaults( + data: Dict[str, Any], configuration: Optional[Config] = None +) -> Defaults: return Defaults( hostname=data.get("hostname"), port=data.get("port"), @@ -42,11 +45,16 @@ def _get_defaults(data: Dict[str, Any]) -> Defaults: platform=data.get("platform"), data=data.get("data"), connection_options=_get_connection_options(data.get("connection_options", {})), + configuration=configuration, ) def _get_inventory_element( - typ: Type[HostOrGroup], data: Dict[str, Any], name: str, defaults: Defaults + typ: Type[HostOrGroup], + data: Dict[str, Any], + name: str, + defaults: Defaults, + configuration: Optional[Config] = None, ) -> HostOrGroup: return typ( name=name, @@ -61,6 +69,7 @@ def _get_inventory_element( ), # this is a hack, we will convert it later to the correct type defaults=defaults, connection_options=_get_connection_options(data.get("connection_options", {})), + configuration=configuration, ) @@ -71,6 +80,7 @@ def __init__( group_file: str = "groups.yaml", defaults_file: str = "defaults.yaml", encoding: str = "utf-8", + configuration: Optional[Config] = None, ) -> None: """ SimpleInventory is an inventory plugin that loads data from YAML files. @@ -90,6 +100,7 @@ def __init__( self.group_file = pathlib.Path(group_file).expanduser() self.defaults_file = pathlib.Path(defaults_file).expanduser() self.encoding = encoding + self._config = configuration def load(self) -> Inventory: yml = ruamel.yaml.YAML(typ="safe") @@ -97,7 +108,7 @@ def load(self) -> Inventory: if self.defaults_file.exists(): with open(self.defaults_file, "r", encoding=self.encoding) as f: defaults_dict = yml.load(f) or {} - defaults = _get_defaults(defaults_dict) + defaults = _get_defaults(defaults_dict, configuration=self._config) else: defaults = Defaults() @@ -106,7 +117,9 @@ def load(self) -> Inventory: hosts_dict = yml.load(f) for n, h in hosts_dict.items(): - hosts[n] = _get_inventory_element(Host, h, n, defaults) + hosts[n] = _get_inventory_element( + Host, h, n, defaults, configuration=self._config + ) groups = Groups() if self.group_file.exists(): @@ -114,7 +127,9 @@ def load(self) -> Inventory: groups_dict = yml.load(f) or {} for n, g in groups_dict.items(): - groups[n] = _get_inventory_element(Group, g, n, defaults) + groups[n] = _get_inventory_element( + Group, g, n, defaults, configuration=self._config + ) for g in groups.values(): g.groups = ParentGroups([groups[g] for g in g.groups]) diff --git a/nornir/plugins/inventory_data/__init__.py b/nornir/plugins/inventory_data/__init__.py new file mode 100644 index 00000000..aae07b69 --- /dev/null +++ b/nornir/plugins/inventory_data/__init__.py @@ -0,0 +1,3 @@ +from .dictionary import InventoryDataDict + +__all__ = ("InventoryDataDict",) diff --git a/nornir/plugins/inventory_data/dictionary.py b/nornir/plugins/inventory_data/dictionary.py new file mode 100644 index 00000000..6c06e84f --- /dev/null +++ b/nornir/plugins/inventory_data/dictionary.py @@ -0,0 +1,24 @@ +import logging +from typing import Any, Dict, Optional, Union + +from nornir.core.plugins.inventory_data import InventoryData + +logger = logging.getLogger(__name__) + + +class InventoryDataDict: + def __init__(self) -> None: + """ + This method configures the plugin + """ + ... + + def load( + self, data: Optional[Dict[str, Any]] + ) -> Union[Dict[str, Any], InventoryData]: + """ + Returns the object containing the data + """ + if data is None: + return dict() + return data diff --git a/pyproject.toml b/pyproject.toml index 1ffc46fd..cf69bf35 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,9 @@ build-backend = "poetry.masonry.api" [tool.poetry.plugins."nornir.plugins.inventory"] "SimpleInventory" = "nornir.plugins.inventory.simple:SimpleInventory" +[tool.poetry.plugins."nornir.plugins.inventory_data"] +"InventoryDataDict" = "nornir.plugins.inventory_data.dictionary:InventoryDataDict" + [tool.poetry] name = "nornir" version = "3.4.1" diff --git a/tests/core/test_configuration.py b/tests/core/test_configuration.py index 9726a912..58c382a9 100644 --- a/tests/core/test_configuration.py +++ b/tests/core/test_configuration.py @@ -25,6 +25,10 @@ def test_config_defaults(self): "transform_function": "", "transform_function_options": {}, }, + "inventory_data": { + "plugin": "InventoryDataDict", + "options": {}, + }, "ssh": {"config_file": str(Path("~/.ssh/config").expanduser())}, "logging": { "enabled": True, @@ -48,6 +52,10 @@ def test_config_from_dict_defaults(self): "transform_function": "", "transform_function_options": {}, }, + "inventory_data": { + "plugin": "InventoryDataDict", + "options": {}, + }, "ssh": {"config_file": str(Path("~/.ssh/config").expanduser())}, "logging": { "enabled": True, @@ -63,6 +71,7 @@ def test_config_from_dict_defaults(self): def test_config_basic(self): c = Config.from_dict( inventory={"plugin": "an-inventory"}, + inventory_data={"plugin": "an-inventory_data"}, runner={"plugin": "serial", "options": {"a": 1, "b": 2}}, logging={"log_file": ""}, user_defined={"my_opt": True}, @@ -74,6 +83,10 @@ def test_config_basic(self): "transform_function": "", "transform_function_options": {}, }, + "inventory_data": { + "plugin": "an-inventory_data", + "options": {}, + }, "runner": {"options": {"a": 1, "b": 2}, "plugin": "serial"}, "ssh": {"config_file": str(Path("~/.ssh/config").expanduser())}, "logging": { diff --git a/tests/core/test_inventory_data.py b/tests/core/test_inventory_data.py new file mode 100644 index 00000000..f1afe057 --- /dev/null +++ b/tests/core/test_inventory_data.py @@ -0,0 +1,188 @@ +# import os + +# import pytest +# import ruamel.yaml + +from typing import Any, Dict, ItemsView, KeysView, ValuesView + +from nornir.core import inventory +from nornir.core.configuration import Config, InventoryDataConfig +from nornir.core.plugins.inventory_data import InventoryDataPluginRegister + +# yaml = ruamel.yaml.YAML(typ="safe") +# dir_path = os.path.dirname(os.path.realpath(__file__)) +# with open(f"{dir_path}/../inventory_data/hosts.yaml") as f: +# hosts = yaml.load(f) +# with open(f"{dir_path}/../inventory_data/groups.yaml") as f: +# groups = yaml.load(f) +# with open(f"{dir_path}/../inventory_data/defaults.yaml") as f: +# defaults = yaml.load(f) +# inv_dict = {"hosts": hosts, "groups": groups, "defaults": defaults} + + +class MockInventoryData: + def __init__(self, data: Dict[str, Any]) -> None: + self.data = data + + def __getitem__(self, key) -> Any: + """ + This method configures the plugin + """ + return self.data[key] + + def get(self, key, default=None) -> Any: + """ + This method configures the plugin + """ + return self.data.get(key, default) + + def __setitem__(self, key, value): + """ + This method configures the plugin + """ + self.data[key] = value + + def keys(self) -> KeysView: + """ + This method configures the plugin + """ + return self.data.keys() + + def values(self) -> ValuesView: + """ + This method configures the plugin + """ + return self.data.values() + + def items(self) -> ItemsView: + """ + This method configures the plugin + """ + return self.data.items() + + +class MockInventoryDataPlugin: + def __init__(self, *args: Any, **kwargs: Any) -> None: + """ + This method configures the plugin + """ + ... + + def load(self, data: Dict[str, Any]) -> Dict[str, Any]: + """ + Returns the object containing the data + """ + return MockInventoryData(data=data) + + +class Test(object): + @classmethod + def setup_class(cls): + InventoryDataPluginRegister.deregister_all() + InventoryDataPluginRegister.register( + "MockInventoryDataPlugin", MockInventoryDataPlugin + ) + + def test_host(self): + config = Config( + inventory_data=InventoryDataConfig(plugin="MockInventoryDataPlugin") + ) + h = inventory.Host( + name="host1", + hostname="host1", + data={"test": "test123"}, + configuration=config, + ) + assert h.hostname == "host1" + assert isinstance(h.data, MockInventoryData) + assert h.data["test"] == "test123" + assert h.get("test") == "test123" + + def test_group(self): + config = Config( + inventory_data=InventoryDataConfig(plugin="MockInventoryDataPlugin") + ) + g = inventory.Group( + name="group1", data={"test": "test123"}, configuration=config + ) + assert g.name == "group1" + assert isinstance(g.data, MockInventoryData) + assert g.data["test"] == "test123" + assert g.data.get("test") == "test123" + + def test_default(self): + config = Config( + inventory_data=InventoryDataConfig(plugin="MockInventoryDataPlugin") + ) + d = inventory.Defaults(data={"test": "test123"}, configuration=config) + assert isinstance(d.data, MockInventoryData) + assert d.data["test"] == "test123" + assert d.data.get("test") == "test123" + + def test_inventory_data(self): + config = Config( + inventory_data=InventoryDataConfig(plugin="MockInventoryDataPlugin") + ) + g1 = inventory.Group( + name="g1", data={"g1": "group1 data"}, configuration=config + ) + g2 = inventory.Group( + name="g2", + groups=inventory.ParentGroups([g1]), + data={"g2": "group2 data"}, + configuration=config, + ) + h1 = inventory.Host( + name="h1", + groups=inventory.ParentGroups([g1, g2]), + data={"host_data": "host 1"}, + configuration=config, + ) + h2 = inventory.Host( + name="h2", data={"host_data": "host 2"}, configuration=config + ) + hosts = {"h1": h1, "h2": h2} + groups = {"g1": g1, "g2": g2} + inv = inventory.Inventory(hosts=hosts, groups=groups) + + assert list(inv.hosts["h2"].keys()) == ["host_data"] + assert "host_data" in inv.hosts["h1"].keys() + assert "g1" in inv.hosts["h1"].keys() + assert "g2" in inv.hosts["h1"].keys() + assert len(inv.hosts["h1"].items()) == 3 + assert len(inv.hosts["h1"].values()) == 3 + + assert inv.hosts["h1"].get("host_data") == "host 1" + assert inv.hosts["h1"].get("g1") == "group1 data" + assert inv.hosts["h1"].get("g2") == "group2 data" + + assert inv.hosts["h1"].extended_data() == { + "host_data": "host 1", + "g1": "group1 data", + "g2": "group2 data", + } + assert inv.hosts["h2"].extended_data() == { + "host_data": "host 2", + } + + def test_inventory_data_default(self): + config = Config( + inventory_data=InventoryDataConfig(plugin="MockInventoryDataPlugin") + ) + d = inventory.Defaults(data={"default": "default data"}, configuration=config) + h1 = inventory.Host( + name="h1", data={"host_data": "host 1"}, defaults=d, configuration=config + ) + hosts = {"h1": h1} + inv = inventory.Inventory(hosts=hosts, defaults=d) + + assert inv.hosts["h1"].get("host_data") == "host 1" + assert inv.hosts["h1"].get("default") == "default data" + assert len(inv.hosts["h1"].keys()) == 2 + assert len(inv.hosts["h1"].items()) == 2 + assert len(inv.hosts["h1"].values()) == 2 + + assert inv.hosts["h1"].extended_data() == { + "host_data": "host 1", + "default": "default data", + }