Merge pull request #660 from AnswerDotAI/args_kwargs

Support args kwargs

Author: jph00

fastcore/docments.py

@@ -65,13 +65,19 @@ def _clean_comment(s):
     res = _clean_re.findall(s)
     return res[0] if res else None
 
-def _param_locs(s, returns=True):
+def _param_locs(s, returns=True, args_kwargs=False):
     "`dict` of parameter line numbers to names"
     body = _parses(s).body
-    if len(body)==1: #or not isinstance(body[0], FunctionDef): return None
+    if len(body)==1:
         defn = body[0]
         if isinstance(defn, (FunctionDef, AsyncFunctionDef)):
             res = {arg.lineno:arg.arg for arg in defn.args.args}
+            # Add *args if present
+            if defn.args.vararg and args_kwargs: res[defn.args.vararg.lineno] = defn.args.vararg.arg
+            # Add keyword-only args
+            if args_kwargs: res.update({arg.lineno:arg.arg for arg in defn.args.kwonlyargs})
+            # Add **kwargs if present
+            if defn.args.kwarg and args_kwargs: res[defn.args.kwarg.lineno] = defn.args.kwarg.arg
             if returns and defn.returns: res[defn.returns.lineno] = 'return'
             return res
         elif isdataclass(s):
@@ -133,12 +139,12 @@ def qual_name(obj):
     return get_name(obj)
 
 # %% ../nbs/06_docments.ipynb
-def _docments(s, returns=True, eval_str=False):
+def _docments(s, returns=True, eval_str=False, args_kwargs=False):
     "`dict` of parameter names to 'docment-style' comments in function or string `s`"
     nps = parse_docstring(s)
     if isclass(s) and not is_dataclass(s): s = s.__init__ # Constructor for a class
     comments = {o.start[0]:_clean_comment(o.string) for o in _tokens(s) if o.type==COMMENT}
-    parms = _param_locs(s, returns=returns) or {}
+    parms = _param_locs(s, returns=returns, args_kwargs=args_kwargs) or {}
     docs = {arg:_get_comment(line, arg, comments, parms) for line,arg in parms.items()}
 
     sig = signature_ex(s, True)

nbs/06_docments.ipynb

@@ -237,13 +237,19 @@
     "    res = _clean_re.findall(s)\n",
     "    return res[0] if res else None\n",
     "\n",
-    "def _param_locs(s, returns=True):\n",
+    "def _param_locs(s, returns=True, args_kwargs=False):\n",
     "    \"`dict` of parameter line numbers to names\"\n",
     "    body = _parses(s).body\n",
-    "    if len(body)==1: #or not isinstance(body[0], FunctionDef): return None\n",
+    "    if len(body)==1:\n",
     "        defn = body[0]\n",
     "        if isinstance(defn, (FunctionDef, AsyncFunctionDef)):\n",
     "            res = {arg.lineno:arg.arg for arg in defn.args.args}\n",
+    "            # Add *args if present\n",
+    "            if defn.args.vararg and args_kwargs: res[defn.args.vararg.lineno] = defn.args.vararg.arg\n",
+    "            # Add keyword-only args\n",
+    "            if args_kwargs: res.update({arg.lineno:arg.arg for arg in defn.args.kwonlyargs})\n",
+    "            # Add **kwargs if present\n",
+    "            if defn.args.kwarg and args_kwargs: res[defn.args.kwarg.lineno] = defn.args.kwarg.arg\n",
     "            if returns and defn.returns: res[defn.returns.lineno] = 'return'\n",
     "            return res\n",
     "        elif isdataclass(s):\n",
@@ -380,12 +386,12 @@
    "outputs": [],
    "source": [
     "#|export\n",
-    "def _docments(s, returns=True, eval_str=False):\n",
+    "def _docments(s, returns=True, eval_str=False, args_kwargs=False):\n",
     "    \"`dict` of parameter names to 'docment-style' comments in function or string `s`\"\n",
     "    nps = parse_docstring(s)\n",
     "    if isclass(s) and not is_dataclass(s): s = s.__init__ # Constructor for a class\n",
     "    comments = {o.start[0]:_clean_comment(o.string) for o in _tokens(s) if o.type==COMMENT}\n",
-    "    parms = _param_locs(s, returns=returns) or {}\n",
+    "    parms = _param_locs(s, returns=returns, args_kwargs=args_kwargs) or {}\n",
     "    docs = {arg:_get_comment(line, arg, comments, parms) for line,arg in parms.items()}\n",
     "\n",
     "    sig = signature_ex(s, True)\n",
@@ -466,6 +472,47 @@
     "docments(add)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/markdown": [
+       "```json\n",
+       "{ 'a': 'the 1st number to add',\n",
+       "  'args': 'some args',\n",
+       "  'b': 'the 2nd number to add',\n",
+       "  'kwargs': 'Passed to the `example` function',\n",
+       "  'return': 'the result of adding `a` to `b`'}\n",
+       "```"
+      ],
+      "text/plain": [
+       "{'args': 'some args',\n",
+       " 'a': 'the 1st number to add',\n",
+       " 'b': 'the 2nd number to add',\n",
+       " 'kwargs': 'Passed to the `example` function',\n",
+       " 'return': 'the result of adding `a` to `b`'}"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "def add(*args, # some args\n",
+    "    a:int, # the 1st number to add\n",
+    "    b=0,   # the 2nd number to add\n",
+    "    **kwargs, # Passed to the `example` function\n",
+    ")->int:    # the result of adding `a` to `b`\n",
+    "    \"The sum of two numbers.\"\n",
+    "    return a+b\n",
+    "\n",
+    "docments(add, args_kwargs=True)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -910,6 +957,37 @@
     "docments(_b)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/markdown": [
+       "```json\n",
+       "{'a': 'First', 'b': 'Second', 'return': None}\n",
+       "```"
+      ],
+      "text/plain": [
+       "{'a': 'First', 'return': None, 'b': 'Second'}"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "def _a(a:int=2): return a # First\n",
+    "\n",
+    "@delegates(_a)\n",
+    "def _b(b:str, # Second\n",
+    "       **kwargs): \n",
+    "    return b, (_a(**kwargs)) \n",
+    "docments(_b)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -1062,13 +1140,6 @@
     "#|hide\n",
     "import nbdev; nbdev.nbdev_export()"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {