Split docname into KnownImport and replace

Reduce responsibility of the former DocName class. Replacing docstring
specific type description should be handled separately.

Author: lagru

src/docstub/_docstrings.py

@@ -1,9 +1,11 @@
 """Transform types defined in docstrings to Python parsable types."""
 
 import logging
+import textwrap
 from dataclasses import dataclass, field
 from pathlib import Path
 
+import click
 import lark
 import lark.visitors
 from numpydoc.docscrape import NumpyDocString
@@ -33,7 +35,7 @@ def _find_one_token(tree: lark.Tree, *, name: str) -> lark.Token:
     return tokens[0]
 
 
-@dataclass(frozen=True, slots=True)
+@dataclass(frozen=True, slots=True, kw_only=True)
 class Annotation:
     """Python-ready type annotation with attached import information."""
 
@@ -148,6 +150,30 @@ def __init__(self, *, inspector, replace_doctypes, **kwargs):
         self._collected_imports = None
         super().__init__(**kwargs)
 
+    def transform(self, doctype):
+        """Turn a type description in a docstring into a type annotation.
+
+        Parameters
+        ----------
+        doctype : str
+            The doctype to parse.
+
+        Returns
+        -------
+        annotation : Annotation
+            The parsed annotation.
+        """
+        try:
+            self._collected_imports = set()
+            tree = _lark.parse(doctype)
+            value = super().transform(tree=tree)
+            annotation = Annotation(
+                value=value, imports=frozenset(self._collected_imports)
+            )
+            return annotation
+        finally:
+            self._collected_imports = None
+
     def __default__(self, data, children, meta):
         """Unpack children of rule nodes by default.
 
@@ -172,30 +198,6 @@ def __default__(self, data, children, meta):
             out = children
         return out
 
-    def transform(self, tree):
-        """
-
-        Parameters
-        ----------
-        tree : lark.Tree
-            The
-
-        Returns
-        -------
-        annotation : Annotation
-            The doctype formatted as a stub-file compatible string with
-            necessary imports attached.
-        """
-        try:
-            self._collected_imports = set()
-            value = super().transform(tree=tree)
-            annotation = Annotation(
-                value=value, imports=frozenset(self._collected_imports)
-            )
-            return annotation
-        finally:
-            self._collected_imports = None
-
     def annotation(self, tree):
         out = " | ".join(tree.children)
         return out
@@ -293,10 +295,20 @@ def _find_import(self, qualname):
 
 
 class DocstringAnnotations:
-    def __init__(self, docstring, *, transformer):
+    def __init__(self, docstring, *, transformer, source=None):
         self.docstring = docstring
         self.np_docstring = NumpyDocString(docstring)
         self.transformer = transformer
+        self.source = source
+
+    def _format_grammar_error(self, error, doctype):
+        msg = "doctype doesn't conform to grammar"
+        details = doctype
+        if hasattr(error, "get_context"):
+            details = error.get_context(doctype)
+        details = textwrap.indent(details, prefix="  ")
+        out = f"{click.style(self.source, bold=True)} {msg}\n{details}"
+        return out
 
     def _doctype_to_annotation(self, doctype):
         """Convert a type description to a Python-ready type.
@@ -316,14 +328,19 @@ def _doctype_to_annotation(self, doctype):
             necessary imports attached.
         """
         try:
-            tree = _lark.parse(doctype)
-            annotation = self.transformer.transform(tree)
+            annotation = self.transformer.transform(doctype)
             return annotation
-        except lark.visitors.VisitError as e:
-            logger.exception("couldn't parse doctype: %r", doctype, exc_info=e.orig_exc)
+        except (lark.exceptions.LexError, lark.exceptions.ParseError) as error:
+            msg = self._format_grammar_error(error=error, doctype=doctype)
+            click.echo(msg)
             return ErrorFallbackAnnotation
-        except Exception:
-            logger.exception("couldn't parse doctype: %r", doctype)
+        except lark.visitors.VisitError as e:
+            logger.exception(
+                "unexpected error parsing doctype %r in %s, falling back to Any",
+                doctype,
+                self.source,
+                exc_info=e.orig_exc,
+            )
             return ErrorFallbackAnnotation
 
     @property

src/docstub/_stubs.py

@@ -191,6 +191,8 @@ class Py2StubTransformer(cst.CSTTransformer):
     inspector : ~._analysis.StaticInspector
     """
 
+    METADATA_DEPENDENCIES = (cst.metadata.PositionProvider,)
+
     # Equivalent to ` ...`, to replace the body of callables with
     _body_replacement = cst.SimpleStatementSuite(
         leading_whitespace=cst.SimpleWhitespace(value=" "),
@@ -240,6 +242,7 @@ def python_to_stub(self, source, *, module_path=None):
                 self.inspector.current_source = module_path
 
             source_tree = cst.parse_module(source)
+            source_tree = cst.metadata.MetadataWrapper(source_tree)
             stub_tree = source_tree.visit(self)
             stub = stub_tree.code
             stub = try_format_stub(stub)
@@ -262,7 +265,7 @@ def visit_ClassDef(self, node):
         out : Literal[True]
         """
         self._scope_stack.append(_Scope(type=FuncType.CLASS, node=node))
-        pytypes = self._pytypes_from_node(node)
+        pytypes = self._annotations_from_node(node)
         self._pytypes_stack.append(pytypes)
         return True
 
@@ -295,7 +298,7 @@ def visit_FunctionDef(self, node):
         """
         func_type = self._function_type(node)
         self._scope_stack.append(_Scope(type=func_type, node=node))
-        pytypes = self._pytypes_from_node(node)
+        pytypes = self._annotations_from_node(node)
         self._pytypes_stack.append(pytypes)
         return True
 
@@ -436,7 +439,7 @@ def visit_Module(self, node):
         Literal[True]
         """
         self._scope_stack.append(_Scope(type=FuncType.MODULE, node=node))
-        pytypes = self._pytypes_from_node(node)
+        pytypes = self._annotations_from_node(node)
         self._pytypes_stack.append(pytypes)
         return True
 
@@ -551,7 +554,7 @@ def _function_type(self, func_def):
                     break
         return func_type
 
-    def _pytypes_from_node(self, node):
+    def _annotations_from_node(self, node):
         """Extract types from function, class or module docstrings.
 
         Parameters
@@ -560,20 +563,21 @@ def _pytypes_from_node(self, node):
 
         Returns
         -------
-        pytypes : dict[str, ~._docstrings.PyType]
+        annotations : DocstringAnnotations
         """
-        pytypes = None
+        annotations = None
         docstring = node.get_docstring()
         if docstring:
+            position = self.get_metadata(cst.metadata.PositionProvider, node).start
+            source = f"{self.inspector.current_source}:{position.line}"
             try:
-                pytypes = DocstringAnnotations(
-                    docstring,
-                    transformer=self.transformer,
+                annotations = DocstringAnnotations(
+                    docstring, transformer=self.transformer, source=source
                 )
             except Exception as e:
                 logger.exception(
                     "error while parsing docstring of `%s`:\n\n%s",
                     node.name.value,
                     e,
                 )
-        return pytypes
+        return annotations