cphyc · May 28, 2021 11:16 · May 28, 2021 · May 28, 2021 · May 28, 2021
diff --git a/check_signature.py b/check_signature.py
@@ -1,28 +1,34 @@
+# %%
 from typing import Callable
 import ast
 from docstring_parser import parse
 from pathlib import Path
 
-def check_non_existing_params(dp_names: list[str], sp_names: list[str], ctx: str):
+# %%
+
+def check_non_existing_params(dp_names: list[str], sp_names: list[str], *, has_args, has_kwargs, ctx: str):
     for dp in dp_names:
-        if dp not in sp_names:
+        if dp not in sp_names and not (has_args or has_kwargs):
             print(f"E100 {ctx}: Argument '{dp}' found in docstring does not exist in function signature.")
 
-def check_params_not_in_docstring(dp_names: list[str], sp_names: list[str], ctx: str):
+def check_params_not_in_docstring(dp_names: list[str], sp_names: list[str], *, has_args, has_kwargs, ctx: str):
     for sp in sp_names:
-        if sp not in dp_names:
+        if sp not in dp_names and not (has_args or has_kwargs):
             print(f"E101 {ctx}: Argument '{sp}' found in signature but not in docstring.")
 
 
-def compare_args(docstring_params, src_params, ctx: str) -> bool:
+def compare_args(docstring_params, src_params, *, has_args, has_kwargs, ctx: str) -> bool:
     if src_params[0].arg in ("self", "cls"):
         src_params = src_params[1:]
 
     # Check ensemble of arguments
-    dp_names = [dp.arg_name for dp in docstring_params]
+    dp_names = []
+    for dp in docstring_params:
+        dp_names.extend(_.strip() for _ in dp.arg_name.replace("*", "").split(","))
+
     sp_names = [sp.arg for sp in src_params]
-    check_non_existing_params(dp_names, sp_names, ctx)
-    # check_params_not_in_docstring(dp_names, sp_names, ctx)
+    check_non_existing_params(dp_names, sp_names, has_args=has_args, has_kwargs=has_kwargs, ctx=ctx)
+    # check_params_not_in_docstring(dp_names, sp_names, has_args=has_args, has_kwargs=has_kwargs, ctx=ctx)
 
 
 def walk_ast_helper(path: Path) -> str:
@@ -35,7 +41,7 @@ def walk_ast_helper(path: Path) -> str:
     nodes = [
         node
         for node in ast.walk(tree)
-        if isinstance(node, (ast.FunctionDef,))
+        if isinstance(node, (ast.ClassDef, ast.FunctionDef))
     ]
 
     # Iterate over docstrings in reversed order so that lines
@@ -46,17 +52,33 @@ def walk_ast_helper(path: Path) -> str:
         docstring = ast.get_docstring(node)
         if not docstring:
             continue
-
         doc = parse(docstring)
-        args = node.args.posonlyargs + node.args.args + node.args.kwonlyargs
-        if not doc.params:
+
+        # For classes, we need to find the __init__ function down there
+        args = []
+        if isinstance(node, ast.ClassDef):
+            for child_node in node.body:
+                if isinstance(child_node, ast.FunctionDef) and child_node.name == "__init__":
+                    args = child_node.args.posonlyargs + child_node.args.args + child_node.args.kwonlyargs
+                    has_kwargs = True if child_node.args.kwarg else False
+                    has_args = True if child_node.args.vararg else False
+                    break
+        else:
+            args = node.args.posonlyargs + node.args.args + node.args.kwonlyargs
+            has_kwargs = True if node.args.kwarg else False
+            has_args = True if node.args.vararg else False
+
+        if not doc.params or not args:
             continue
         ctx = f"{path}:{node.lineno}"
-        compare_args(doc.params, args, ctx=ctx)
+        compare_args(doc.params, args, has_kwargs=has_kwargs, has_args=has_args, ctx=ctx)
 
 
+# %%
 files = list((Path(__file__).parent / ".." ).glob("**/*.py"))
 print(f"Found {len(files)} files")
 
 for f in files:
-    walk_ast_helper(f)
+    walk_ast_helper(f)
+
+# %%
diff --git a/check_signature → check_signature.py b/check_signature → check_signature.py
@@ -1,14 +1,7 @@
-# %%
 from typing import Callable
 import ast
 from docstring_parser import parse
 from pathlib import Path
-import yt
-
-ds = yt.load_sample("output_00080")
-p = yt.ProjectionPlot(ds, "x", "density")
-
-# %%
 
 def check_non_existing_params(dp_names: list[str], sp_names: list[str], ctx: str):
     for dp in dp_names:
@@ -62,11 +55,8 @@ def walk_ast_helper(path: Path) -> str:
         compare_args(doc.params, args, ctx=ctx)
 
 
-# %%
 files = list((Path(__file__).parent / ".." ).glob("**/*.py"))
 print(f"Found {len(files)} files")
 
 for f in files:
-    walk_ast_helper(f)
-
-# %%
+    walk_ast_helper(f)
diff --git a/check_signature b/check_signature
@@ -0,0 +1,72 @@
+# %%
+from typing import Callable
+import ast
+from docstring_parser import parse
+from pathlib import Path
+import yt
+
+ds = yt.load_sample("output_00080")
+p = yt.ProjectionPlot(ds, "x", "density")
+
+# %%
+
+def check_non_existing_params(dp_names: list[str], sp_names: list[str], ctx: str):
+    for dp in dp_names:
+        if dp not in sp_names:
+            print(f"E100 {ctx}: Argument '{dp}' found in docstring does not exist in function signature.")
+
+def check_params_not_in_docstring(dp_names: list[str], sp_names: list[str], ctx: str):
+    for sp in sp_names:
+        if sp not in dp_names:
+            print(f"E101 {ctx}: Argument '{sp}' found in signature but not in docstring.")
+
+
+def compare_args(docstring_params, src_params, ctx: str) -> bool:
+    if src_params[0].arg in ("self", "cls"):
+        src_params = src_params[1:]
+
+    # Check ensemble of arguments
+    dp_names = [dp.arg_name for dp in docstring_params]
+    sp_names = [sp.arg for sp in src_params]
+    check_non_existing_params(dp_names, sp_names, ctx)
+    # check_params_not_in_docstring(dp_names, sp_names, ctx)
+
+
+def walk_ast_helper(path: Path) -> str:
+    src = path.read_text()
+    lines = src.splitlines()
+    newLines = lines.copy()
+
+    # Extract all functions and classes
+    tree = ast.parse(src)
+    nodes = [
+        node
+        for node in ast.walk(tree)
+        if isinstance(node, (ast.FunctionDef,))
+    ]
+
+    # Iterate over docstrings in reversed order so that lines
+    # can be modified
+    for node in sorted(
+        nodes, key=lambda node: node.body[0].lineno if node.body else 0
+    ):
+        docstring = ast.get_docstring(node)
+        if not docstring:
+            continue
+
+        doc = parse(docstring)
+        args = node.args.posonlyargs + node.args.args + node.args.kwonlyargs
+        if not doc.params:
+            continue
+        ctx = f"{path}:{node.lineno}"
+        compare_args(doc.params, args, ctx=ctx)
+
+
+# %%
+files = list((Path(__file__).parent / ".." ).glob("**/*.py"))
+print(f"Found {len(files)} files")
+
+for f in files:
+    walk_ast_helper(f)
+
+# %%
No results found