From f550907543d4e327c7385499743fcb9ae07d9f00 Mon Sep 17 00:00:00 2001
From: Kefu Chai <kchai@redhat.com>
Date: Sat, 26 Dec 2020 20:26:55 +0800
Subject: [PATCH] pybing/mgr/mgr_module: use annotations when collecting args

Signed-off-by: Kefu Chai <kchai@redhat.com>
---
 src/mypy.ini                 |  6 ++++-
 src/pybind/mgr/mgr_module.py | 48 +++++++++++++++++++++++++++++++++---
 src/pybind/mgr/tox.ini       |  5 ++++
 3 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/src/mypy.ini b/src/mypy.ini
index 53b2a5cfd94e7..3371557f10d7b 100755
--- a/src/mypy.ini
+++ b/src/mypy.ini
@@ -64,4 +64,8 @@ ignore_missing_imports = True
 
 # Make volumes happy:
 [mypy-StringIO]
-ignore_missing_imports = True
\ No newline at end of file
+ignore_missing_imports = True
+
+[mypy-ceph_argparse]
+# more work to do
+ignore_errors = True
diff --git a/src/pybind/mgr/mgr_module.py b/src/pybind/mgr/mgr_module.py
index b207a3c676a35..9b8235d8cdc92 100644
--- a/src/pybind/mgr/mgr_module.py
+++ b/src/pybind/mgr/mgr_module.py
@@ -1,6 +1,6 @@
 import ceph_module  # noqa
 
-from typing import Set, Tuple, Iterator, Any, Dict, Optional, Callable, List, \
+from typing import Set, Tuple, Iterator, Any, Dict, Generic, Optional, Callable, List, \
     Union, TYPE_CHECKING, NamedTuple
 if TYPE_CHECKING:
     import sys
@@ -9,7 +9,6 @@ if TYPE_CHECKING:
     else:
         from typing_extensions import Literal
 
-
 import inspect
 import logging
 import errno
@@ -19,9 +18,24 @@ import threading
 from collections import defaultdict, namedtuple
 import rados
 import re
+import sys
 import time
+from ceph_argparse import CephArgtype
 from mgr_util import profile_method
 
+if sys.version_info >= (3, 8):
+    from typing import get_args, get_origin
+else:
+    def get_args(tp):
+        if tp is Generic:
+            return tp
+        else:
+            return getattr(tp, '__args__', ())
+
+    def get_origin(tp):
+        return getattr(tp, '__origin__', None)
+
+
 ERROR_MSG_EMPTY_INPUT_FILE = 'Empty content: please add a password/secret to the file.'
 ERROR_MSG_NO_INPUT_FILE = 'Please specify the file containing the password/secret with "-i" option.'
 # Full list of strings in "osd_types.cc:pg_state_string()"
@@ -294,6 +308,8 @@ class CLICommand(object):
         self.func = None  # type: Optional[Callable]
         self.args = args
         self.args_dict = self._parse_args(args)
+        self.arg_spec = {}    # type: Dict[str, Any]
+        self.first_default = -1
 
     @staticmethod
     def _parse_args(args):
@@ -315,9 +331,14 @@ class CLICommand(object):
         self.func = func
         if not self.desc:
             self.desc = inspect.getdoc(func)
+        if not self.args_dict:
+            self.arg_spec = inspect.getfullargspec(func).annotations
         self.COMMANDS[self.prefix] = self
         return self.func
 
+    def _is_arg_key(self, k):
+        return k in self.args_dict or k in self.arg_spec
+
     def _get_arg_value(self, kwargs_switch, key, val):
         def start_kwargs():
             if isinstance(val, str) and '=' in val:
@@ -335,7 +356,7 @@ class CLICommand(object):
             k, v = key, val
         return kwargs_switch, k.replace('-', '_'), v
 
-    def _collect_args(self, cmd_dict):
+    def _collect_args_by_argdesc(self, cmd_dict):
         kwargs = {}
         kwargs_switch = False
         for a, d in self.args_dict.items():
@@ -346,8 +367,27 @@ class CLICommand(object):
             kwargs[k] = v
         return kwargs
 
+    def _collect_args_by_argspec(self, cmd_dict):
+        kwargs = {}
+        kwargs_switch = False
+        for index, (name, tp) in enumerate(self.arg_spec.items()):
+            if name in CLICommand.KNOWN_ARGS:
+                continue
+            assert self.first_default >= 0
+            raw_v = cmd_dict.get(name)
+            if index >= self.first_default:
+                if raw_v is None:
+                    continue
+            kwargs_switch, k, v = self._get_arg_value(kwargs_switch,
+                                                      name, raw_v)
+            kwargs[k] = v
+        return kwargs
+
     def call(self, mgr, cmd_dict, inbuf):
-        kwargs = self._collect_args(cmd_dict)
+        if self.args_dict:
+            kwargs = self._collect_args_by_argdesc(cmd_dict)
+        else:
+            kwargs = self._collect_args_by_argspec(cmd_dict)
         if inbuf:
             kwargs['inbuf'] = inbuf
         assert self.func
diff --git a/src/pybind/mgr/tox.ini b/src/pybind/mgr/tox.ini
index bb9c62e880064..020930e27390c 100644
--- a/src/pybind/mgr/tox.ini
+++ b/src/pybind/mgr/tox.ini
@@ -28,6 +28,7 @@ addopts =
 [testenv]
 setenv =
     UNITTEST = true
+    PYTHONPATH = $PYTHONPATH:..
 deps =
     cython
     -rrequirements.txt
@@ -43,6 +44,10 @@ commands =
         snap_schedule}
 
 [testenv:mypy]
+setenv =
+    MYPYPATH = {toxinidir}/..
+passenv =
+    MYPYPATH
 basepython = python3
 deps =
     cython
-- 
2.39.5