Stub-related improvements (#386)

TypeAlias, TypedDict, eval functions. And runtime TypeAlias.

Author: JelleZijlstra

docs/changelog.md

@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+- Support some imports from stub-only modules (#386)
+- Support type evaluation functions in stubs (#386)
+- Support `TypedDict` in stubs (#386)
+- Support `TypeAlias` (PEP 612) (#386)
 - Small improvements to `ParamSpec` support (#385)
 - Allow `CustomCheck` to customize what values
   a value can be assigned to (#383)

pyanalyze/annotations.py

@@ -30,7 +30,9 @@
 import qcore
 import ast
 import builtins
+import inspect
 from collections.abc import Callable, Iterable
+import textwrap
 from typing import (
     Any,
     Container,
@@ -46,7 +48,7 @@
     Union,
     TYPE_CHECKING,
 )
-from typing_extensions import ParamSpec
+from typing_extensions import ParamSpec, TypedDict
 
 from .error_code import ErrorCode
 from .extensions import (
@@ -57,8 +59,11 @@
     NoReturnGuard,
     ParameterTypeGuard,
     TypeGuard,
+    get_type_evaluation,
 )
 from .find_unused import used
+from .functions import FunctionDefNode
+from .node_visitor import ErrorContext
 from .signature import ELLIPSIS_PARAM, SigParameter, Signature, ParameterKind
 from .safe import is_typing_name, is_instance_of_typing_name
 from . import type_evaluation
@@ -67,7 +72,6 @@
     AnySource,
     AnyValue,
     CallableValue,
-    CanAssignContext,
     CustomCheckExtension,
     Extension,
     HasAttrGuardExtension,
@@ -152,11 +156,9 @@ def get_name_from_globals(self, name: str, globals: Mapping[str, Any]) -> Value:
 
 
 @dataclass
-class TypeEvaluationContext(Context, type_evaluation.Context):
-    variables: type_evaluation.VarMap
-    positions: Mapping[str, type_evaluation.Position]
-    can_assign_context: CanAssignContext
+class RuntimeEvaluator(type_evaluation.Evaluator, Context):
     globals: Mapping[str, object] = field(repr=False)
+    func: typing.Callable[..., Any]
 
     def evaluate_type(self, node: ast.AST) -> Value:
         return type_from_ast(node, ctx=self)
@@ -168,36 +170,58 @@ def get_name(self, node: ast.Name) -> Value:
         """Return the :class:`Value <pyanalyze.value.Value>` corresponding to a name."""
         return self.get_name_from_globals(node.id, self.globals)
 
+    @classmethod
+    def get_for(cls, func: typing.Callable[..., Any]) -> Optional["RuntimeEvaluator"]:
+        try:
+            key = f"{func.__module__}.{func.__qualname__}"
+        except AttributeError:
+            return None
+        evaluation_func = get_type_evaluation(key)
+        if evaluation_func is None or not hasattr(evaluation_func, "__globals__"):
+            return None
+        lines, _ = inspect.getsourcelines(evaluation_func)
+        code = textwrap.dedent("".join(lines))
+        body = ast.parse(code)
+        if not body.body:
+            return None
+        evaluator = body.body[0]
+        if not isinstance(evaluator, ast.FunctionDef):
+            return None
+        return RuntimeEvaluator(evaluator, evaluation_func.__globals__, evaluation_func)
+
 
 @dataclass
-class EvaluatorValidationContext(Context, type_evaluation.Context):
-    variables: type_evaluation.VarMap
-    positions: Mapping[str, type_evaluation.Position]
-    can_assign_context: "NameCheckVisitor"
-    node: ast.AST
+class SyntheticEvaluator(type_evaluation.Evaluator):
+    error_ctx: ErrorContext
+    annotations_context: Context
 
     def show_error(
         self,
         message: str,
         error_code: ErrorCode = ErrorCode.invalid_annotation,
         node: Optional[ast.AST] = None,
     ) -> None:
-        self.can_assign_context.show_error(
-            node or self.node, message, error_code=error_code
-        )
+        self.error_ctx.show_error(node or self.node, message, error_code=error_code)
 
     def evaluate_type(self, node: ast.AST) -> Value:
-        return type_from_ast(node, ctx=self)
+        return type_from_ast(node, ctx=self.annotations_context)
 
     def evaluate_value(self, node: ast.AST) -> Value:
-        return value_from_ast(node, ctx=self, error_on_unrecognized=False)
+        return value_from_ast(
+            node, ctx=self.annotations_context, error_on_unrecognized=False
+        )
 
     def get_name(self, node: ast.Name) -> Value:
         """Return the :class:`Value <pyanalyze.value.Value>` corresponding to a name."""
-        val, _ = self.can_assign_context.resolve_name(
-            node, suppress_errors=self.should_suppress_undefined_names
+        return self.annotations_context.get_name(node)
+
+    @classmethod
+    def from_visitor(
+        cls, node: FunctionDefNode, visitor: "NameCheckVisitor"
+    ) -> "SyntheticEvaluator":
+        return cls(
+            node, visitor, _DefaultContext(visitor, node, use_name_node_for_error=True)
         )
-        return val
 
 
 @used  # part of an API
@@ -223,6 +247,25 @@ def type_from_ast(
     return _type_from_ast(ast_node, ctx)
 
 
+def type_from_annotations(
+    annotations: Mapping[str, object],
+    key: str,
+    *,
+    globals: Optional[Mapping[str, object]] = None,
+    ctx: Optional[Context] = None,
+) -> Optional[Value]:
+    try:
+        annotation = annotations[key]
+    except Exception:
+        # Malformed __annotations__
+        return None
+    else:
+        maybe_val = type_from_runtime(annotation, globals=globals, ctx=ctx)
+        if maybe_val != AnyValue(AnySource.incomplete_annotation):
+            return maybe_val
+    return None
+
+
 def type_from_runtime(
     val: object,
     visitor: Optional["NameCheckVisitor"] = None,
@@ -490,11 +533,15 @@ def _type_from_runtime(val: Any, ctx: Context, is_typeddict: bool = False) -> Va
     elif is_instance_of_typing_name(val, "_MaybeRequired"):
         required = is_instance_of_typing_name(val, "_Required")
         if is_typeddict:
-            return _Pep655Value(required, _type_from_runtime(val.__type__, ctx))
+            return Pep655Value(required, _type_from_runtime(val.__type__, ctx))
         else:
             cls = "Required" if required else "NotRequired"
             ctx.show_error(f"{cls}[] used in unsupported context")
             return AnyValue(AnySource.error)
+    elif is_typing_name(val, "TypeAlias"):
+        return AnyValue(AnySource.incomplete_annotation)
+    elif is_typing_name(val, "TypedDict"):
+        return KnownValue(TypedDict)
     else:
         origin = get_origin(val)
         if isinstance(origin, type):
@@ -571,7 +618,7 @@ def _get_typeddict_value(
     total: bool,
 ) -> Tuple[bool, Value]:
     val = _type_from_runtime(value, ctx, is_typeddict=True)
-    if isinstance(val, _Pep655Value):
+    if isinstance(val, Pep655Value):
         return (val.required, val.value)
     if required_keys is None:
         required = total
@@ -605,6 +652,8 @@ def _type_from_value(value: Value, ctx: Context, is_typeddict: bool = False) ->
         )
     elif isinstance(value, AnyValue):
         return value
+    elif isinstance(value, SubclassValue) and value.exactly:
+        return value.typ
     elif isinstance(value, TypedValue) and isinstance(value.typ, str):
         # Synthetic type
         return value
@@ -634,6 +683,15 @@ def _type_from_subscripted_value(
                 for subval in root.vals
             ]
         )
+    if (
+        isinstance(root, SubclassValue)
+        and root.exactly
+        and isinstance(root.typ, TypedValue)
+    ):
+        return GenericValue(
+            root.typ.typ, [_type_from_value(elt, ctx) for elt in members]
+        )
+
     if isinstance(root, TypedValue) and isinstance(root.typ, str):
         return GenericValue(root.typ, [_type_from_value(elt, ctx) for elt in members])
 
@@ -691,15 +749,15 @@ def _type_from_subscripted_value(
         if len(members) != 1:
             ctx.show_error("Required[] requires a single argument")
             return AnyValue(AnySource.error)
-        return _Pep655Value(True, _type_from_value(members[0], ctx))
+        return Pep655Value(True, _type_from_value(members[0], ctx))
     elif is_typing_name(root, "NotRequired"):
         if not is_typeddict:
             ctx.show_error("NotRequired[] used in unsupported context")
             return AnyValue(AnySource.error)
         if len(members) != 1:
             ctx.show_error("NotRequired[] requires a single argument")
             return AnyValue(AnySource.error)
-        return _Pep655Value(False, _type_from_value(members[0], ctx))
+        return Pep655Value(False, _type_from_value(members[0], ctx))
     elif root is Callable or root is typing.Callable:
         if len(members) == 2:
             args, return_value = members
@@ -739,11 +797,13 @@ def __init__(
         visitor: "NameCheckVisitor",
         node: Optional[ast.AST],
         globals: Optional[Mapping[str, object]] = None,
+        use_name_node_for_error: bool = False,
     ) -> None:
         super().__init__()
         self.visitor = visitor
         self.node = node
         self.globals = globals
+        self.use_name_node_for_error = use_name_node_for_error
 
     def show_error(
         self,
@@ -760,7 +820,7 @@ def get_name(self, node: ast.Name) -> Value:
         if self.visitor is not None:
             val, _ = self.visitor.resolve_name(
                 node,
-                error_node=self.node,
+                error_node=node if self.use_name_node_for_error else self.node,
                 suppress_errors=self.should_suppress_undefined_names,
             )
             return val
@@ -786,7 +846,7 @@ class _SubscriptedValue(Value):
 
 
 @dataclass
-class _Pep655Value(Value):
+class Pep655Value(Value):
     required: bool
     value: Value
 
@@ -901,7 +961,6 @@ def visit_Call(self, node: ast.Call) -> Optional[Value]:
                     if isinstance(kwarg_value, KnownValue):
                         kwargs[name] = kwarg_value.val
                     else:
-                        print(kwarg_value)
                         return None
             return KnownValue(func.val(*args, **kwargs))
         elif func.val == TypeVar:
@@ -1038,15 +1097,15 @@ def _value_of_origin_args(
         if len(args) != 1:
             ctx.show_error("Required[] requires a single argument")
             return AnyValue(AnySource.error)
-        return _Pep655Value(True, _type_from_runtime(args[0], ctx))
+        return Pep655Value(True, _type_from_runtime(args[0], ctx))
     elif is_typing_name(origin, "NotRequired"):
         if not is_typeddict:
             ctx.show_error("NotRequired[] used in unsupported context")
             return AnyValue(AnySource.error)
         if len(args) != 1:
             ctx.show_error("NotRequired[] requires a single argument")
             return AnyValue(AnySource.error)
-        return _Pep655Value(False, _type_from_runtime(args[0], ctx))
+        return Pep655Value(False, _type_from_runtime(args[0], ctx))
     elif origin is None and isinstance(val, type):
         # This happens for SupportsInt in 3.7.
         return _maybe_typed_value(val)

pyanalyze/arg_spec.py

@@ -7,7 +7,7 @@
 from .options import Options, PyObjectSequenceOption
 from .analysis_lib import is_positional_only_arg_name
 from .extensions import CustomCheck, get_overloads
-from .annotations import Context, type_from_runtime
+from .annotations import Context, RuntimeEvaluator, type_from_runtime
 from .config import Config
 from .find_unused import used
 from . import implementation
@@ -33,7 +33,6 @@
     Signature,
     ParameterKind,
 )
-from .type_evaluation import get_evaluator
 from .typeshed import TypeshedFinder
 from .value import (
     AnySource,
@@ -94,9 +93,9 @@ def _scribe_log_impl(variables, visitor, node):
         ):
             yield
     else:
-        argspec = ArgSpecCache(Options.from_option_list([], Config())).get_argspec(
-            fn, impl=implementation_fn
-        )
+        options = Options.from_option_list([], Config())
+        tsf = TypeshedFinder.make(options)
+        argspec = ArgSpecCache(options, tsf).get_argspec(fn, impl=implementation_fn)
         if argspec is None:
             # builtin or something, just use a generic argspec
             argspec = Signature.make(
@@ -288,13 +287,14 @@ class ArgSpecCache:
     def __init__(
         self,
         options: Options,
+        ts_finder: TypeshedFinder,
         *,
         vnv_provider: Callable[[str], Optional[Value]] = lambda _: None,
     ) -> None:
         self.vnv_provider = vnv_provider
         self.options = options
         self.config = options.fallback
-        self.ts_finder = TypeshedFinder(verbose=False)
+        self.ts_finder = ts_finder
         self.known_argspecs = {}
         self.generic_bases_cache = {}
         self.default_context = AnnotationsContext(self)
@@ -509,7 +509,7 @@ def _uncached_get_argspec(
                     ]
                     if all_of_type(sigs, Signature):
                         return OverloadedSignature(sigs)
-                evaluator = get_evaluator(obj)
+                evaluator = RuntimeEvaluator.get_for(obj)
                 if evaluator is not None:
                     sig = self._cached_get_argspec(
                         evaluator.func, impl, is_asynq, in_overload_resolution=True

pyanalyze/attributes.py

@@ -14,7 +14,7 @@
 from typing import Any, Generic, Sequence, Tuple, Optional, Union
 
 
-from .annotations import type_from_runtime, Context
+from .annotations import type_from_annotations, type_from_runtime, Context
 from .safe import safe_isinstance, safe_issubclass
 from .signature import Signature, MaybeSignature
 from .stacked_scopes import Composite
@@ -376,13 +376,14 @@ def _get_attribute_from_mro(
             pass
     elif safe_isinstance(typ, types.ModuleType):
         try:
-            annotation = typ.__annotations__[ctx.attr]
+            annotations = typ.__annotations__
         except Exception:
-            # Module doesn't have annotations or it's not in there
             pass
         else:
-            attr_type = type_from_runtime(annotation, ctx=AnnotationsContext(ctx, typ))
-            if attr_type != AnyValue(AnySource.incomplete_annotation):
+            attr_type = type_from_annotations(
+                annotations, ctx.attr, ctx=AnnotationsContext(ctx, typ)
+            )
+            if attr_type is not None:
                 return (attr_type, typ, False)
 
     try:
@@ -407,15 +408,15 @@ def _get_attribute_from_mro(
             try:
                 # Make sure to use only __annotations__ that are actually on this
                 # class, not ones inherited from a base class.
-                annotation = base_cls.__dict__["__annotations__"][ctx.attr]
+                annotations = base_cls.__dict__["__annotations__"]
             except Exception:
                 # no __annotations__, or it's not a dict, or the attr isn't there
                 pass
             else:
-                attr_type = type_from_runtime(
-                    annotation, ctx=AnnotationsContext(ctx, base_cls)
+                attr_type = type_from_annotations(
+                    annotations, ctx.attr, ctx=AnnotationsContext(ctx, base_cls)
                 )
-                if attr_type != AnyValue(AnySource.incomplete_annotation):
+                if attr_type is not None:
                     return (attr_type, base_cls, False)
             try:
                 # Make sure we use only the object from this class, but do invoke