ENH: add runtests option to find untyped module attributes

Now you can run something like ``` ./runtests.py --find-missing numpy ``` to get a list of everything that still needs annotations added.
numpy · shoyer · Apr 22, 2020 · Apr 16, 2020 · Apr 18, 2020 · Apr 22, 2020
commit b7224d93e7c8071d9fdb89b1ef9015bf94562255
diff --git a/runtests.py b/runtests.py
@@ -1,4 +1,7 @@
 #!/usr/bin/env python
+import argparse
+import ast
+import importlib
 import os
 import subprocess
 import sys
@@ -8,13 +11,84 @@
 STUBS_ROOT = os.path.dirname(os.path.abspath(__file__))
 
 
-def main():
+class FindAttributes(ast.NodeVisitor):
+    """Find top-level attributes/functions/classes in the stubs.
+
+    Do this by walking the stubs ast. See e.g.
+
+    https://greentreesnakes.readthedocs.io/en/latest/index.html
+
+    for more information on working with Python's ast.
+
+    """
+
+    def __init__(self):
+        self.attributes = set()
+
+    def visit_FunctionDef(self, node):
+        if node.name == '__getattr__':
+            # Not really a module member.
+            return
+        self.attributes.add(node.name)
+        # Do not call self.generic_visit; we are only interested in
+        # top-level functions.
+        return
+
+    def visit_ClassDef(self, node):
+        if not node.name.startswith('_'):
+            self.attributes.add(node.name)
+        return
+
+    def visit_AnnAssign(self, node):
+        self.attributes.add(node.target.id)
+
+
+def find_missing(module_name):
+    module_path = os.path.join(
+        STUBS_ROOT,
+        module_name.replace('numpy', 'numpy-stubs').replace('.', os.sep),
+        '__init__.pyi',
+    )
+
+    module = importlib.import_module(module_name)
+    module_attributes = {
+        attribute for attribute in dir(module)
+        if not attribute.startswith('_')
+    }
+
+    if os.path.isfile(module_path):
+        with open(module_path) as f:
+            tree = ast.parse(f.read())
+        ast_visitor = FindAttributes()
+        ast_visitor.visit(tree)
+        stubs_attributes = ast_visitor.attributes
+    else:
+        # No stubs for this module yet.
+        stubs_attributes = set()
+
+    missing = module_attributes - stubs_attributes
+    print('\n'.join(sorted(missing)))
+
+
+def run_pytest(argv):
     subprocess.run(
         [sys.executable, "-m", "pip", "install", STUBS_ROOT],
         capture_output=True,
         check=True,
     )
-    sys.exit(pytest.main([STUBS_ROOT] + sys.argv[1:]))
+    return pytest.main([STUBS_ROOT] + argv)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--find-missing')
+    args, remaining_argv = parser.parse_known_args()
+
+    if args.find_missing is not None:
+        find_missing(args.find_missing)
+        sys.exit(0)
+
+    sys.exit(run_pytest(remaining_argv))
 
 
 if __name__ == "__main__":