Merge pull request #8 from timeforplanb123/updates

Add full support for json-strings for options and arguments

.gitignore

@@ -7,6 +7,7 @@ dist/
 # Service files
 config.yaml
 nornir.log
+temp.pkl
 .temp.pkl
 
 *.swp

examples/custom_runbooks/aaa/cmd_create_fw_users.py

@@ -29,9 +29,7 @@
 SPECIAL_CHARACTERS = "'\"[!@#$%^&*()-+?_=,<>}{~:]/\"'"
 
 
-def validate_user_parameters(
-    task, valid_user_parameters, current_user_parameters
-):
+def validate_user_parameters(task, valid_user_parameters, current_user_parameters):
     for dict_0, dict_1 in zip(valid_user_parameters, current_user_parameters):
         if dict_0 != dict_1:
             return Result(host=task.host, result="Validation error")
@@ -119,9 +117,7 @@ def create_fw_users(task):
         template = task.run(
             task=template_file,
             name="Create template",
-            path=os.path.join(
-                os.path.dirname(os.path.abspath(__file__)), "templates"
-            ),
+            path=os.path.join(os.path.dirname(os.path.abspath(__file__)), "templates"),
             template="create_fw_user.j2",
         )
 
@@ -165,15 +161,13 @@ def create_fw_users(task):
         f, ch = (result[host].failed, result[host].changed)
         if f:
             click.secho(
-                f"{host:<25}: Oh! Here is some exception:"
-                f" {result[host].exception}",
+                f"{host:<25}: Oh! Here is some exception:" f" {result[host].exception}",
                 fg="red",
                 bold=True,
             )
         elif ch:
             click.secho(
-                f"{host:<25}: {' '.join(usernames)}"
-                " have been created or changed",
+                f"{host:<25}: {' '.join(usernames)}" " have been created or changed",
                 fg="yellow",
                 bold=True,
             )

examples/custom_runbooks/dhcp_snooping/cmd_dhcp_snooping.py

@@ -28,7 +28,8 @@ def _get_trusted_untrusted(task):
             command_string="disp int",
             use_textfsm=True,
             textfsm_template=os.path.join(
-                os.path.dirname(os.path.abspath(__file__)), "templates/disp_int.template"
+                os.path.dirname(os.path.abspath(__file__)),
+                "templates/disp_int.template",
             ),
         )
         # Get trusted interfaces
@@ -44,9 +45,7 @@ def _get_trusted_untrusted(task):
         # Render j2 template
         template = task.run(
             task=template_file,
-            path=os.path.join(
-                os.path.dirname(os.path.abspath(__file__)), "templates"
-            ),
+            path=os.path.join(os.path.dirname(os.path.abspath(__file__)), "templates"),
             template="dhcp_snooping.j2",
         )
         # Configure commands from j2 template

nornir_cli/common_commands/__init__.py

@@ -1,6 +1,7 @@
 from nornir_cli.common_commands.common import (
     _doc_generator,
     _get_lists,
+    _get_dict_from_json_string,
     _json_loads,
     _pickle_to_hidden_file,
     _validate_connection_options,
@@ -27,6 +28,7 @@
 __all__ = (
     "_doc_generator",
     "_get_lists",
+    "_get_dict_from_json_string",
     "_json_loads",
     "_pickle_to_hidden_file",
     "_validate_connection_options",

nornir_cli/common_commands/cmd_filter.py

@@ -7,6 +7,7 @@
 from nornir_cli.common_commands import (
     SHOW_INVENTORY_OPTIONS,
     _get_lists,
+    _get_dict_from_json_string,
     _json_loads,
     _pickle_to_hidden_file,
     common_options,
@@ -90,17 +91,15 @@ def cli(ctx, advanced_filter, f, save, **kwargs):
         if advanced_filter:
             ctx.obj["nornir"] = _get_obj_after_adv_filter(ctx.obj["nornir"], f)
         else:
-            d = dict(
-                [_json_loads(i) for i in (value.split("=") for value in _get_lists(f))]
-            )
+            d = _get_dict_from_json_string(f)
             ctx.obj["nornir"] = ctx.obj["nornir"].filter(**d)
         if save:
             _pickle_to_hidden_file("temp.pkl", obj=ctx.obj["nornir"])
 
         # run show_inventory command
         if any(kwargs.values()):
             ctx.invoke(show_inventory, **kwargs)
-    except (ValueError, IndexError):
+    except (ValueError, TypeError, IndexError):
         raise ctx.fail(
             click.BadParameter(
                 f"{ERROR_MESSAGE}",

nornir_cli/common_commands/cmd_init.py

@@ -7,6 +7,7 @@
     CONNECTION_OPTIONS,
     SHOW_INVENTORY_OPTIONS,
     _get_lists,
+    _get_dict_from_json_string,
     _json_loads,
     _pickle_to_hidden_file,
     common_options,
@@ -93,16 +94,10 @@ def cli(
     ctx.ensure_object(dict)
 
     try:
-
         TransformFunctionRegister.register("adapt_host_data", adapt_host_data)
 
         if from_dict:
-            d = dict(
-                [
-                    _json_loads(i)
-                    for i in (value.split("=") for value in _get_lists(from_dict))
-                ]
-            )
+            d = _get_dict_from_json_string(from_dict)
             cf = (
                 None
                 if not config_file or config_file == "None" or "null"
@@ -135,11 +130,11 @@ def cli(
                 f"{ERROR_MESSAGE}",
             ).format_message(),
         )
-    except (AttributeError):
+    except AttributeError:
         ctx.fail(
             f"File '{config_file}' is empty",
         )
-    except (FileNotFoundError):
+    except FileNotFoundError:
         raise ctx.fail(
             click.BadParameter(
                 f"Path '{config_file}' does not exist",

nornir_cli/common_commands/cmd_show_inventory.py

@@ -11,7 +11,6 @@
 
 
 def _get_inventory(nr_obj, count, **kwargs):
-
     d = {
         str: dict,
         bool: list,
@@ -24,7 +23,6 @@ def _get_inventory(nr_obj, count, **kwargs):
 
     for k, v in kwargs.items():
         if v:
-
             if v == "all":
                 for inventory_key in ("hosts", "groups", "defaults"):
                     for item in _get_inventory(nr_obj, count, inventory=inventory_key):

nornir_cli/common_commands/cmd_write_file.py

@@ -18,7 +18,6 @@ def cli(ctx, filename, content, append, line_feed):
     """
 
     try:
-
         dirname = os.path.dirname(filename)
         Path(dirname).mkdir(parents=True, exist_ok=True)
 

nornir_cli/common_commands/common.py

@@ -139,7 +139,7 @@ def _get_lists(s):
         begin = list(takewhile(lambda x: "=" in x, s))
         s = list(dropwhile(lambda x: "=" in x, s))
         begin.extend([i for i in takewhile(lambda x: "=" not in x, s)])
-        body.append("".join(begin))
+        body.append(" ".join(begin))
         s = list(dropwhile(lambda x: "=" not in x, s))
     return body
 
@@ -209,3 +209,16 @@ def multiple_progress_bar(task, method, pg_bar, **kwargs):
     task.run(task=method, **kwargs)
     if pg_bar:
         pg_bar.update()
+
+
+# json_string to dict function
+def _get_dict_from_json_string(json_string):
+    if "=" in json_string:
+        return dict(
+            [
+                _json_loads(i)
+                for i in (value.split("=") for value in _get_lists(json_string))
+            ]
+        )
+    else:
+        return _json_loads([json_string])[0]

nornir_cli/common_commands/static_options.py

@@ -36,6 +36,10 @@
         default=True,
         show_default=True,
     ),
+    click.argument(
+        "arguments",
+        required=False,
+    ),
 ]
 
 PRINT_RESULT_WRITE_RESULT_OPTIONS = [

nornir_cli/common_commands/write_result.py

@@ -45,7 +45,6 @@ def _write_individual_result(
     write_host: bool = False,
     no_errors: bool = False,
 ) -> None:
-
     # ignore results with a specifig severity_level
     if result.severity_level < severity_level:
         return
@@ -93,7 +92,6 @@ def _write_result(
     count: Optional[int] = None,
     no_errors: bool = False,
 ) -> None:
-
     attrs = attrs or ["diff", "result", "stdout"]
     if isinstance(attrs, str):
         attrs = [attrs]

nornir_cli/common_commands/write_results.py

@@ -52,7 +52,6 @@ def _write_individual_result(
     no_errors: bool = False,
     append: bool = False,
 ) -> None:
-
     individual_result = []
 
     # ignore results with a specifig severity_level
@@ -120,13 +119,11 @@ def _write_results(
     no_errors: bool = False,
     append: bool = False,
 ) -> None:
-
     attrs = attrs or ["diff", "result", "stdout"]
     if isinstance(attrs, str):
         attrs = [attrs]
 
     if isinstance(result, AggregatedResult):
-
         result = dict(sorted(result.items()))
 
         if isinstance(count, int):

nornir_cli/nornir_cli.py

@@ -195,7 +195,6 @@ def finder():
                         for filename in p[2]
                         if filename.endswith(".py") and filename.startswith("cmd_")
                     ]:
-
                         cmd_path = p[0].split(path)[1]
 
                         grps = cmd_path.split(os.sep)[1:]
@@ -305,16 +304,22 @@ def wrapper(f):
                 str(default_value) if not isinstance(default_value, type) else None
             )
 
+            if not default_value and default_value not in ["", 0]:
+                ctx.obj["required_options"].append(k)
+
             click.option(
                 "--" + k,
                 default=default_value,
                 show_default=True,
-                required=False if default_value or default_value == "" else True,
+                required=False,
+                help="[required]" if k in ctx.obj["required_options"] else "",
                 type=PARAMETER_TYPES.setdefault(typ, click.STRING),
             )(cmd)
 
             # last original functions with arguments
-            ctx.obj["queue_parameters"][obj_or].update({k: v.default})
+            ctx.obj["queue_parameters"][obj_or].update(
+                {k: v.default if not isinstance(v.default, type) else default_value}
+            )
 
         # list of dictionaries with original function (key) and set of arguments (value)
         ctx.obj["queue_functions"].append(ctx.obj["queue_parameters"])
@@ -339,6 +344,9 @@ def wrapper(f):
 
     ctx.obj["queue_parameters"][obj_or] = {}
 
+    # list with required options for obj_or
+    ctx.obj["required_options"] = []
+
     return wrapper
 
 
@@ -357,6 +365,8 @@ def init_nornir_cli(ctx):
     ctx.obj["queue_functions"] = []
     # last original functions with arguments
     ctx.obj["queue_parameters"] = {}
+    # list with required options for last plugin command
+    ctx.obj["required_options"] = []
 
 
 @dec("nornir_netmiko")

nornir_cli/plugin_commands/cmd_common.py

@@ -3,6 +3,8 @@
 from nornir.core.plugins.connections import ConnectionPluginRegister
 
 from nornir_cli.common_commands import (
+    _get_lists,
+    _get_dict_from_json_string,
     _json_loads,
     _pickle_to_hidden_file,
     multiple_progress_bar,
@@ -13,24 +15,49 @@
 from tqdm import tqdm
 
 
+ERROR_MESSAGE = "error"
+
+
 @click.pass_context
-def cli(ctx, pg_bar, print_result, print_stat, *args, **kwargs):
+def cli(ctx, pg_bar, print_result, print_stat, arguments, *args, **kwargs):
     ConnectionPluginRegister.auto_register()
 
-    # 'None' = None
-    kwargs.update({k: v for k, v in zip(kwargs, _json_loads(kwargs.values()))})
+    try:
+        # 'None' = None
+        kwargs.update({k: v for k, v in zip(kwargs, _json_loads(kwargs.values()))})
 
-    current_func_params = next(ctx.obj["queue_functions_generator"])
+        current_func_params = next(ctx.obj["queue_functions_generator"])
 
-    # try to pass not all arguments, but only the necessary ones
-    if kwargs == list(current_func_params.values())[0]:
-        function, parameters = list(current_func_params.items())[0]
-    else:
-        param_iterator = iter(current_func_params.values())
-        params = list(next(param_iterator).items())
-        function, parameters = list(current_func_params)[0], {
-            key: value for key, value in kwargs.items() if (key, value) not in params
-        }
+        # try to pass not all arguments, but only the necessary ones
+        if kwargs == list(current_func_params.values())[0] and not arguments:
+            function, parameters = list(current_func_params.items())[0]
+        else:
+            param_iterator = iter(current_func_params.values())
+            params = list(next(param_iterator).items())
+            function, parameters = list(current_func_params)[0], {
+                key: value
+                for key, value in kwargs.items()
+                if (key, value) not in params
+            }
+
+            # use arguments to update parameters
+            if arguments:
+                parameters = {**_get_dict_from_json_string(arguments), **parameters}
+        missing_options = [
+            f"'--{option}'"
+            for option in ctx.obj["required_options"]
+            if parameters.get(option) == None
+        ]
+        if missing_options:
+            raise ctx.fail(
+                f"Missing options {', '.join(missing_options)}",
+            )
+    except (ValueError, IndexError, TypeError, KeyError):
+        raise ctx.fail(
+            click.BadParameter(
+                f"{ERROR_MESSAGE}",
+            ).format_message(),
+        )
 
     # get the current Nornir object from Commands chain or from temp.pkl file
     try: