Skip to content

Commit 3bd33e4

Browse files
authored
Support attributes and type aliases (#18)
* Support collection of explicit type aliases While it is very difficult to statically detect a type alias [1] such as `x = str`, explicit ones are easier to detect. This change makes sure that `x: TypeAlias = str` and `type x = str` are collected as importable type annotations. [1] https://docs.python.org/3/library/typing.html#type-aliases * Handle attributes in stub generation Type information in docstrings is not yet taken into account. Instead, `typing.Any` is used as a fallback. * Use `_typeshed.Incomplete` as a fallback type instead * Refactor enum `FuncType` to `ScopeType` * Add support for attribute annotation with doctypes as well as proper handling in context of stubs. In stubs, the value should basically always be omitted from attributes except for a few exceptions. * Handle attributes in stub generation Type information in docstrings is not yet taken into account. Instead, `typing.Any` is used as a fallback. * Blacklist reserved Python keywords from qualnames
1 parent ef8de11 commit 3bd33e4

File tree

7 files changed

+548
-57
lines changed

7 files changed

+548
-57
lines changed

examples/example_pkg-stubs/_basic.pyi

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,11 @@ import logging
44
from collections.abc import Sequence
55
from typing import Any, Literal, Self, Union
66

7+
from _typeshed import Incomplete
8+
79
from . import CustomException
810

9-
logger = ...
11+
logger: Incomplete
1012

1113
__all__ = [
1214
"func_empty",

src/docstub/_analysis.py

Lines changed: 60 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,11 @@
77
import re
88
import typing
99
from dataclasses import asdict, dataclass
10+
from functools import cache
1011
from pathlib import Path
1112

1213
import libcst as cst
14+
import libcst.matchers as cstm
1315

1416
from ._utils import accumulate_qualname, module_name_from_path, pyfile_checksum
1517

@@ -45,13 +47,13 @@ class KnownImport:
4547
4648
Attributes
4749
----------
48-
import_path : str, optional
50+
import_path
4951
Dotted names after "from".
50-
import_name : str, optional
52+
import_name
5153
Dotted names after "import".
52-
import_alias : str, optional
54+
import_alias
5355
Name (without ".") after "as".
54-
builtin_name : str, optional
56+
builtin_name
5557
Names an object that's builtin and doesn't need an import.
5658
5759
Examples
@@ -65,6 +67,26 @@ class KnownImport:
6567
import_alias: str = None
6668
builtin_name: str = None
6769

70+
@classmethod
71+
@cache
72+
def typeshed_Incomplete(cls):
73+
"""Create import corresponding to ``from _typeshed import Incomplete``.
74+
75+
This type is not actually available at runtime and only intended to be
76+
used in stub files [1]_.
77+
78+
Returns
79+
-------
80+
import : KnownImport
81+
The import corresponding to ``from _typeshed import Incomplete``.
82+
83+
References
84+
----------
85+
.. [1] https://typing.readthedocs.io/en/latest/guides/writing_stubs.html#incomplete-stubs
86+
"""
87+
import_ = cls(import_path="_typeshed", import_name="Incomplete")
88+
return import_
89+
6890
@classmethod
6991
def one_from_config(cls, name, *, info):
7092
"""Create one KnownImport from the configuration format.
@@ -327,23 +349,47 @@ def __init__(self, *, module_name):
327349

328350
def visit_ClassDef(self, node: cst.ClassDef) -> bool:
329351
self._stack.append(node.name.value)
330-
331-
class_name = ".".join(self._stack[:1])
332-
qualname = f"{self.module_name}.{'.'.join(self._stack)}"
333-
known_import = KnownImport(import_path=self.module_name, import_name=class_name)
334-
self.known_imports[qualname] = known_import
335-
352+
self._collect_type_annotation(self._stack)
336353
return True
337354

338355
def leave_ClassDef(self, original_node: cst.ClassDef) -> None:
339356
self._stack.pop()
340357

341358
def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
342-
self._stack.append(node.name.value)
343-
return True
359+
return False
360+
361+
def visit_TypeAlias(self, node: cst.TypeAlias) -> bool:
362+
"""Collect type alias with 3.12 syntax."""
363+
stack = [*self._stack, node.name.value]
364+
self._collect_type_annotation(stack)
365+
return False
366+
367+
def visit_AnnAssign(self, node: cst.AnnAssign) -> bool:
368+
"""Collect type alias annotated with `TypeAlias`."""
369+
is_type_alias = cstm.matches(
370+
node,
371+
cstm.AnnAssign(
372+
annotation=cstm.Annotation(annotation=cstm.Name(value="TypeAlias"))
373+
),
374+
)
375+
if is_type_alias and node.value is not None:
376+
names = cstm.findall(node.target, cstm.Name())
377+
assert len(names) == 1
378+
stack = [*self._stack, names[0].value]
379+
self._collect_type_annotation(stack)
380+
return False
344381

345-
def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None:
346-
self._stack.pop()
382+
def _collect_type_annotation(self, stack):
383+
"""Collect an importable type annotation.
384+
385+
Parameters
386+
----------
387+
stack : Iterable[str]
388+
A list of names that form the path to the collected type.
389+
"""
390+
qualname = ".".join([self.module_name, *stack])
391+
known_import = KnownImport(import_path=self.module_name, import_name=stack[0])
392+
self.known_imports[qualname] = known_import
347393

348394

349395
class TypesDatabase:

src/docstub/_docstrings.py

Lines changed: 90 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
from numpydoc.docscrape import NumpyDocString
1414

1515
from ._analysis import KnownImport
16-
from ._utils import ContextFormatter, accumulate_qualname, escape_qualname
16+
from ._utils import ContextFormatter, DocstubError, accumulate_qualname, escape_qualname
1717

1818
logger = logging.getLogger(__name__)
1919

@@ -135,16 +135,29 @@ def _aggregate_annotations(*types):
135135
return values, imports
136136

137137

138-
GrammarErrorFallback = Annotation(
139-
value="Any",
140-
imports=frozenset((KnownImport(import_path="typing", import_name="Any"),)),
138+
FallbackAnnotation = Annotation(
139+
value="Incomplete", imports=frozenset([KnownImport.typeshed_Incomplete()])
141140
)
142141

143142

143+
class QualnameIsKeyword(DocstubError):
144+
"""Raised when a qualname is a blacklisted Python keyword."""
145+
146+
144147
@lark.visitors.v_args(tree=True)
145148
class DoctypeTransformer(lark.visitors.Transformer):
146149
"""Transformer for docstring type descriptions (doctypes).
147150
151+
Attributes
152+
----------
153+
blacklisted_qualnames : frozenset[str]
154+
All Python keywords [1]_ are blacklisted from use in qualnames except for ``True``
155+
``False`` and ``None``.
156+
157+
References
158+
----------
159+
.. [1] https://docs.python.org/3/reference/lexical_analysis.html#keywords
160+
148161
Examples
149162
--------
150163
>>> transformer = DoctypeTransformer()
@@ -155,6 +168,43 @@ class DoctypeTransformer(lark.visitors.Transformer):
155168
[('tuple', 0, 5), ('int', 9, 12)]
156169
"""
157170

171+
blacklisted_qualnames = frozenset(
172+
{
173+
"await",
174+
"else",
175+
"import",
176+
"pass",
177+
"break",
178+
"except",
179+
"in",
180+
"raise",
181+
"class",
182+
"finally",
183+
"is",
184+
"return",
185+
"and",
186+
"continue",
187+
"for",
188+
"lambda",
189+
"try",
190+
"as",
191+
"def",
192+
"from",
193+
"nonlocal",
194+
"while",
195+
"assert",
196+
"del",
197+
"global",
198+
"not",
199+
"with",
200+
"async",
201+
"elif",
202+
"if",
203+
"or",
204+
"yield",
205+
}
206+
)
207+
158208
def __init__(self, *, types_db=None, replace_doctypes=None, **kwargs):
159209
"""
160210
Parameters
@@ -204,7 +254,11 @@ def doctype_to_annotation(self, doctype):
204254
value=value, imports=frozenset(self._collected_imports)
205255
)
206256
return annotation, self._unknown_qualnames
207-
except (lark.exceptions.LexError, lark.exceptions.ParseError):
257+
except (
258+
lark.exceptions.LexError,
259+
lark.exceptions.ParseError,
260+
QualnameIsKeyword,
261+
):
208262
self.stats["grammar_errors"] += 1
209263
raise
210264
finally:
@@ -274,6 +328,13 @@ def qualname(self, tree):
274328

275329
_qualname = self._find_import(_qualname, meta=tree.meta)
276330

331+
if _qualname in self.blacklisted_qualnames:
332+
msg = (
333+
f"qualname {_qualname!r} in docstring type description "
334+
"is a reserved Python keyword and not allowed"
335+
)
336+
raise QualnameIsKeyword(msg)
337+
277338
_qualname = lark.Token(type="QUALNAME", value=_qualname)
278339
return _qualname
279340

@@ -399,7 +460,7 @@ def _doctype_to_annotation(self, doctype, ds_line=0):
399460
details = details.replace("^", click.style("^", fg="red", bold=True))
400461
if ctx:
401462
ctx.print_message("invalid syntax in doctype", details=details)
402-
return GrammarErrorFallback
463+
return FallbackAnnotation
403464

404465
except lark.visitors.VisitError as e:
405466
tb = "\n".join(traceback.format_exception(e.orig_exc))
@@ -408,7 +469,7 @@ def _doctype_to_annotation(self, doctype, ds_line=0):
408469
ctx.print_message(
409470
"unexpected error while parsing doctype", details=details
410471
)
411-
return GrammarErrorFallback
472+
return FallbackAnnotation
412473

413474
else:
414475
for name, start_col, stop_col in unknown_qualnames:
@@ -421,6 +482,28 @@ def _doctype_to_annotation(self, doctype, ds_line=0):
421482
)
422483
return annotation
423484

485+
@cached_property
486+
def attributes(self) -> dict[str, Annotation]:
487+
annotations = {}
488+
for attribute in self.np_docstring["Attributes"]:
489+
if not attribute.type:
490+
continue
491+
492+
ds_line = 0
493+
for i, line in enumerate(self.docstring.split("\n")):
494+
if attribute.name in line and attribute.type in line:
495+
ds_line = i
496+
break
497+
498+
if attribute.name in annotations:
499+
logger.warning("duplicate parameter name %r, ignoring", attribute.name)
500+
continue
501+
502+
annotation = self._doctype_to_annotation(attribute.type, ds_line=ds_line)
503+
annotations[attribute.name] = annotation
504+
505+
return annotations
506+
424507
@cached_property
425508
def parameters(self) -> dict[str, Annotation]:
426509
all_params = chain(

0 commit comments

Comments
 (0)