Uncouple formatting from Py2StubTransformer

Author: lagru

src/docstub/_cli.py

@@ -14,7 +14,12 @@
 )
 from ._cache import FileCache
 from ._config import Config
-from ._stubs import Py2StubTransformer, walk_source, walk_source_and_targets
+from ._stubs import (
+    Py2StubTransformer,
+    try_format_stub,
+    walk_source,
+    walk_source_and_targets,
+)
 from ._version import __version__
 
 logger = logging.getLogger(__name__)
@@ -171,6 +176,7 @@ def main(source_dir, out_dir, config_path, verbose):
                 stub_content = stub_transformer.python_to_stub(
                     py_content, module_path=source_path
                 )
+                stub_content = try_format_stub(stub_content)
             except (SystemExit, KeyboardInterrupt):
                 raise
             except Exception as e:

src/docstub/_stubs.py

@@ -264,7 +264,7 @@ def current_source(self, value):
         if self.types_db is not None:
             self.types_db.current_source = value
 
-    def python_to_stub(self, source, *, module_path=None, try_format=True):
+    def python_to_stub(self, source, *, module_path=None):
         """Convert Python source code to stub-file ready code.
 
         Parameters
@@ -274,9 +274,6 @@ def python_to_stub(self, source, *, module_path=None, try_format=True):
             The location of the source that is transformed into a stub file.
             If given, used to enhance logging & error messages with more
             context information.
-        try_format : bool, optional
-            Try to format the output, if the appropriate dependencies are
-            installed.
 
         Returns
         -------
@@ -292,8 +289,6 @@ def python_to_stub(self, source, *, module_path=None, try_format=True):
             source_tree = cst.metadata.MetadataWrapper(source_tree)
             stub_tree = source_tree.visit(self)
             stub = stub_tree.code
-            if try_format is True:
-                stub = try_format_stub(stub)
             return stub
         finally:
             self._scope_stack = None

tests/test_stubs.py

@@ -173,7 +173,7 @@ def test_attributes_no_doctype(self, assign, expected, scope):
             src = NESTED_CLASS_ATTRIBUTE_TEMPLATE.format(assign=assign, doctype="")
 
         transformer = Py2StubTransformer()
-        result = transformer.python_to_stub(src, try_format=False)
+        result = transformer.python_to_stub(src)
 
         # Find exactly one occurrence of `expected`
         pattern = f"^ *({re.escape(expected)})$"
@@ -212,7 +212,7 @@ def test_attributes_with_doctype(self, assign, doctype, expected, scope):
             src = NESTED_CLASS_ATTRIBUTE_TEMPLATE.format(assign=assign, doctype=doctype)
 
         transformer = Py2StubTransformer()
-        result = transformer.python_to_stub(src, try_format=False)
+        result = transformer.python_to_stub(src)
 
         # Find exactly one occurrence of `expected`
         pattern = f"^ *({re.escape(expected)})$"