src/phi/cli.py

@@ -2,16 +2,17 @@ import sys
 import argparse
 import inspect
 from phi.logging import setup_logging, get_logger
+from phi.compat.argparse import BooleanOptionalAction
 
 log = get_logger(__name__)
 
 parser = argparse.ArgumentParser()
-subparses = parser.add_subparsers(title='actions', dest='action')
+subparses = parser.add_subparsers(title="actions", dest="action")
 
 cli_callbacks = {}
 
 
-def register(action_info='', param_infos=[]):
+def register(action_info="", param_infos=[]):
     def decorator(action):
         # Get function name and arguments
         action_name = action.__name__
@@ -21,7 +22,7 @@ def register(action_info='', param_infos=[]):
         subparser = subparses.add_parser(action_name, help=action_info)
 
         for i, name in enumerate(param_names):
-            info = param_infos[i] if i<len(param_infos) else ''
+            info = param_infos[i] if i < len(param_infos) else ""
             subparser.add_argument(dest=name, help=info)
 
         # Register action
@@ -33,7 +34,7 @@ def register(action_info='', param_infos=[]):
 
 def run(args):
     for action_name, (action, param_names) in cli_callbacks.items():
-        if args['action'] == action_name:
+        if args["action"] == action_name:
             action(**{pname: args[pname] for pname in param_names})
 
 
@@ -41,6 +42,12 @@ def add_arg(name, example, info):
     parser.add_argument(name, metavar=example, help=info)
 
 
+def add_flag(name, info):
+    parser.add_argument(
+        name, action=BooleanOptionalAction, help=info, required=False, default=False
+    )
+
+
 def get_args():
     namespace = parser.parse_args(sys.argv[1:])
     args = namespace.__dict__

src/phi/compat/argparse.py

@@ -0,0 +1,51 @@
+# -*- encoding: utf-8 -*-
+# WARNING: This code has been copied verbatim from
+# https://github.com/python/cpython/blob/3.9/Lib/argparse.py
+# To be removed when the minimum python supported version
+# will be >= 3.9
+
+import argparse
+
+
+class BooleanOptionalAction(argparse.Action):
+    def __init__(
+        self,
+        option_strings,
+        dest,
+        default=None,
+        type=None,
+        choices=None,
+        required=False,
+        help=None,
+        metavar=None,
+    ):
+
+        _option_strings = []
+        for option_string in option_strings:
+            _option_strings.append(option_string)
+
+            if option_string.startswith("--"):
+                option_string = "--no-" + option_string[2:]
+                _option_strings.append(option_string)
+
+        if help is not None and default is not None:
+            help += f" (default: {default})"
+
+        super().__init__(
+            option_strings=_option_strings,
+            dest=dest,
+            nargs=0,
+            default=default,
+            type=type,
+            choices=choices,
+            required=required,
+            help=help,
+            metavar=metavar,
+        )
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        if option_string in self.option_strings:
+            setattr(namespace, self.dest, not option_string.startswith("--no-"))
+
+    def format_usage(self):
+        return " | ".join(self.option_strings)