Add --include-private option to stubgen to include private members. (#3367)

* Add --include-private option to stubgen to include private members.

* use "flags" instead of "opts" to identify cli flags in stub tests.

* require the list of arguments to be passed to stubgen.parse_options.

Author: chadrik

mypy/stubgen.py

@@ -73,6 +73,7 @@
                                  ('modules', List[str]),
                                  ('ignore_errors', bool),
                                  ('recursive', bool),
+                                 ('include_private', bool),
                                  ])
 
 
@@ -86,7 +87,8 @@ def generate_stub_for_module(module: str, output_dir: str, quiet: bool = False,
                              pyversion: Tuple[int, int] = defaults.PYTHON3_VERSION,
                              no_import: bool = False,
                              search_path: List[str] = [],
-                             interpreter: str = sys.executable) -> None:
+                             interpreter: str = sys.executable,
+                             include_private: bool = False) -> None:
     target = module.replace('.', '/')
     try:
         result = find_module_path_and_all(module=module,
@@ -118,7 +120,7 @@ def generate_stub_for_module(module: str, output_dir: str, quiet: bool = False,
         target = os.path.join(output_dir, target)
         generate_stub(module_path, output_dir, module_all,
                       target=target, add_header=add_header, module=module,
-                      pyversion=pyversion)
+                      pyversion=pyversion, include_private=include_private)
     if not quiet:
         print('Created %s' % target)
 
@@ -185,7 +187,8 @@ def load_python_module_info(module: str, interpreter: str) -> Tuple[str, Optiona
 
 def generate_stub(path: str, output_dir: str, _all_: Optional[List[str]] = None,
                   target: str = None, add_header: bool = False, module: str = None,
-                  pyversion: Tuple[int, int] = defaults.PYTHON3_VERSION
+                  pyversion: Tuple[int, int] = defaults.PYTHON3_VERSION,
+                  include_private: bool = False
                   ) -> None:
     with open(path, 'rb') as f:
         source = f.read()
@@ -199,7 +202,7 @@ def generate_stub(path: str, output_dir: str, _all_: Optional[List[str]] = None,
             sys.stderr.write('%s\n' % m)
         sys.exit(1)
 
-    gen = StubGenerator(_all_, pyversion=pyversion)
+    gen = StubGenerator(_all_, pyversion=pyversion, include_private=include_private)
     ast.accept(gen)
     if not target:
         target = os.path.join(output_dir, os.path.basename(path))
@@ -223,7 +226,8 @@ def generate_stub(path: str, output_dir: str, _all_: Optional[List[str]] = None,
 
 
 class StubGenerator(mypy.traverser.TraverserVisitor):
-    def __init__(self, _all_: Optional[List[str]], pyversion: Tuple[int, int]) -> None:
+    def __init__(self, _all_: Optional[List[str]], pyversion: Tuple[int, int],
+                 include_private: bool = False) -> None:
         self._all_ = _all_
         self._output = []  # type: List[str]
         self._import_lines = []  # type: List[str]
@@ -235,6 +239,7 @@ def __init__(self, _all_: Optional[List[str]], pyversion: Tuple[int, int]) -> No
         self._classes = set()  # type: Set[str]
         self._base_classes = []  # type: List[str]
         self._pyversion = pyversion
+        self._include_private = include_private
 
     def visit_mypy_file(self, o: MypyFile) -> None:
         self._classes = find_classes(o)
@@ -515,6 +520,8 @@ def is_not_in_all(self, name: str) -> bool:
         return False
 
     def is_private_name(self, name: str) -> bool:
+        if self._include_private:
+            return False
         return name.startswith('_') and (not name.endswith('__')
                                          or name in ('__all__',
                                                      '__author__',
@@ -610,7 +617,7 @@ def walk_packages(packages: List[str]) -> Iterator[str]:
 
 
 def main() -> None:
-    options = parse_options()
+    options = parse_options(sys.argv[1:])
     if not os.path.isdir('out'):
         raise SystemExit('Directory "out" does not exist')
     if options.recursive and options.no_import:
@@ -636,23 +643,24 @@ def main() -> None:
                                      pyversion=options.pyversion,
                                      no_import=options.no_import,
                                      search_path=options.search_path,
-                                     interpreter=options.interpreter)
+                                     interpreter=options.interpreter,
+                                     include_private=options.include_private)
         except Exception as e:
             if not options.ignore_errors:
                 raise e
             else:
                 print("Stub generation failed for", module, file=sys.stderr)
 
 
-def parse_options() -> Options:
-    args = sys.argv[1:]
+def parse_options(args: List[str]) -> Options:
     pyversion = defaults.PYTHON3_VERSION
     no_import = False
     recursive = False
     ignore_errors = False
     doc_dir = ''
     search_path = []  # type: List[str]
     interpreter = ''
+    include_private = False
     while args and args[0].startswith('-'):
         if args[0] == '--doc-dir':
             doc_dir = args[1]
@@ -673,6 +681,8 @@ def parse_options() -> Options:
             pyversion = defaults.PYTHON2_VERSION
         elif args[0] == '--no-import':
             no_import = True
+        elif args[0] == '--include-private':
+            include_private = True
         elif args[0] in ('-h', '--help'):
             usage()
         else:
@@ -689,7 +699,8 @@ def parse_options() -> Options:
                    interpreter=interpreter,
                    modules=args,
                    ignore_errors=ignore_errors,
-                   recursive=recursive)
+                   recursive=recursive,
+                   include_private=include_private)
 
 
 def default_python2_interpreter() -> str:
@@ -721,6 +732,9 @@ def usage() -> None:
           --no-import     don't import the modules, just parse and analyze them
                           (doesn't work with C extension modules and doesn't
                           respect __all__)
+          --include-private
+                          generate stubs for objects and members considered private
+                          (single leading undescore and no trailing underscores)
           --doc-dir PATH  use .rst documentation in PATH (this may result in
                           better stubs in some cases; consider setting this to
                           DIR/Python-X.Y.Z/Doc/library)

mypy/test/teststubgen.py

@@ -6,6 +6,7 @@
 import sys
 import tempfile
 import time
+import re
 from types import ModuleType
 
 from typing import List, Tuple
@@ -16,7 +17,7 @@
 from mypy.test import config
 from mypy.parse import parse
 from mypy.errors import CompileError
-from mypy.stubgen import generate_stub, generate_stub_for_module
+from mypy.stubgen import generate_stub, generate_stub_for_module, parse_options, Options
 from mypy.stubgenc import generate_c_type_stub, infer_method_sig
 from mypy.stubutil import (
     parse_signature, parse_all_signatures, build_signature, find_unique_signatures,
@@ -104,11 +105,21 @@ def cases(self) -> List[DataDrivenTestCase]:
         return c
 
 
+def parse_flags(program_text: str) -> Options:
+    flags = re.search('# flags: (.*)$', program_text, flags=re.MULTILINE)
+    if flags:
+        flag_list = flags.group(1).split()
+    else:
+        flag_list = []
+    return parse_options(flag_list + ['dummy.py'])
+
+
 def test_stubgen(testcase: DataDrivenTestCase) -> None:
     if 'stubgen-test-path' not in sys.path:
         sys.path.insert(0, 'stubgen-test-path')
     os.mkdir('stubgen-test-path')
     source = '\n'.join(testcase.input)
+    options = parse_flags(source)
     handle = tempfile.NamedTemporaryFile(prefix='prog_', suffix='.py', dir='stubgen-test-path',
                                          delete=False)
     assert os.path.isabs(handle.name)
@@ -125,9 +136,11 @@ def test_stubgen(testcase: DataDrivenTestCase) -> None:
         reset_importlib_caches()
         try:
             if testcase.name.endswith('_import'):
-                generate_stub_for_module(name, out_dir, quiet=True)
+                generate_stub_for_module(name, out_dir, quiet=True,
+                                         no_import=options.no_import,
+                                         include_private=options.include_private)
             else:
-                generate_stub(path, out_dir)
+                generate_stub(path, out_dir, include_private=options.include_private)
             a = load_output(out_dir)
         except CompileError as e:
             a = e.messages

test-data/unit/stubgen.test

@@ -125,25 +125,52 @@ class A: ...
 [out]
 class A: ...
 
-[case testPrivateFunction]
+[case testSkipPrivateFunction]
 def _f(): ...
 def g(): ...
 [out]
 def g(): ...
 
-[case testPrivateMethod]
+[case testIncludePrivateFunction]
+# flags:  --include-private
+def _f(): ...
+def g(): ...
+[out]
+def _f(): ...
+def g(): ...
+
+[case testSkipPrivateMethod]
 class A:
     def _f(self): ...
 [out]
 class A: ...
 
-[case testPrivateVar]
+[case testIncludePrivateMethod]
+# flags:  --include-private
+class A:
+    def _f(self): ...
+[out]
+class A:
+    def _f(self): ...
+
+[case testSkipPrivateVar]
 _x = 1
 class A:
     _y = 1
 [out]
 class A: ...
 
+[case testIncludePrivateVar]
+# flags:  --include-private
+_x = 1
+class A:
+    _y = 1
+[out]
+_x = ...  # type: int
+
+class A:
+    _y = ...  # type: int
+
 [case testSpecialInternalVar]
 __all__ = []
 __author__ = ''
@@ -309,14 +336,30 @@ class A:
     x = ...  # type: int
     def f(self): ...
 
-[case testMultiplePrivateDefs]
+[case testSkipMultiplePrivateDefs]
+class A: ...
+_x = 1
+_y = 1
+_z = 1
+class C: ...
+[out]
+class A: ...
+class C: ...
+
+[case testIncludeMultiplePrivateDefs]
+# flags:  --include-private
 class A: ...
 _x = 1
 _y = 1
 _z = 1
 class C: ...
 [out]
 class A: ...
+
+_x = ...  # type: int
+_y = ...  # type: int
+_z = ...  # type: int
+
 class C: ...
 
 [case testIncludeFromImportIfInAll_import]
@@ -396,6 +439,16 @@ class A:
 [out]
 class A: ...
 
+[case testIncludePrivateProperty]
+# flags:  --include-private
+class A:
+    @property
+    def _foo(self): ...
+[out]
+class A:
+    @property
+    def _foo(self): ...
+
 [case testSkipPrivateStaticAndClassMethod]
 class A:
     @staticmethod
@@ -405,6 +458,20 @@ class A:
 [out]
 class A: ...
 
+[case testIncludePrivateStaticAndClassMethod]
+# flags:  --include-private
+class A:
+    @staticmethod
+    def _foo(): ...
+    @classmethod
+    def _bar(cls): ...
+[out]
+class A:
+    @staticmethod
+    def _foo(): ...
+    @classmethod
+    def _bar(cls): ...
+
 [case testNamedtuple]
 import collections, x
 X = collections.namedtuple('X', ['a', 'b'])