From 1972669f647d11d179c70c818e237b6eafe0097e Mon Sep 17 00:00:00 2001 From: Isaac Flath Date: Thu, 26 Dec 2024 11:46:35 -0700 Subject: [PATCH] Support args kwargs --- fastcore/docments.py | 14 +++++-- nbs/06_docments.ipynb | 93 ++++++++++++++++++++++++++++++++++++++----- 2 files changed, 92 insertions(+), 15 deletions(-) diff --git a/fastcore/docments.py b/fastcore/docments.py index 8cc33483..02efe9ed 100644 --- a/fastcore/docments.py +++ b/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) diff --git a/nbs/06_docments.ipynb b/nbs/06_docments.ipynb index 01c8f98e..3d1eb751 100644 --- a/nbs/06_docments.ipynb +++ b/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": {