Make _LiteralEvalVisitor private again; expose a new evaluate_expression_truthiness function instead (#130)

AlexWaygood · web-flow · commit dc933baf47b3 · 2025-07-23T14:38:19.000-07:00
diff --git a/typeshed_client/__init__.py b/typeshed_client/__init__.py
@@ -11,10 +11,11 @@
 )
 from .parser import (
     ImportedName,
-    LiteralEvalVisitor,
+    InvalidStub,
     NameDict,
     NameInfo,
     OverloadedName,
+    evaluate_expression_truthiness,
     get_stub_names,
     parse_ast,
 )
@@ -26,14 +27,15 @@
 __all__ = [
     "ImportedInfo",
     "ImportedName",
-    "LiteralEvalVisitor",
+    "InvalidStub",
     "ModulePath",
     "NameDict",
     "NameInfo",
     "OverloadedName",
     "Resolver",
     "SearchContext",
     "__version__",
+    "evaluate_expression_truthiness",
     "get_all_stub_files",
     "get_search_context",
     "get_stub_ast",
diff --git a/typeshed_client/parser.py b/typeshed_client/parser.py
@@ -56,7 +56,6 @@ def get_stub_names(
         ModulePath(tuple(module_name.split("."))),
         is_init=is_init,
         file_path=path,
-        is_py_file=path.suffix == ".py",
     )
 
 
@@ -65,16 +64,11 @@ def parse_ast(
     search_context: SearchContext,
     module_name: ModulePath,
     *,
+    file_path: Path,
     is_init: bool = False,
-    file_path: Optional[Path] = None,
-    is_py_file: bool = False,
 ) -> NameDict:
     visitor = _NameExtractor(
-        search_context,
-        module_name,
-        is_init=is_init,
-        file_path=file_path,
-        is_py_file=is_py_file,
+        search_context, module_name, is_init=is_init, file_path=file_path
     )
     name_dict: NameDict = {}
     try:
@@ -208,15 +202,17 @@ def __init__(
         ctx: SearchContext,
         module_name: ModulePath,
         *,
+        file_path: Path,
         is_init: bool = False,
-        file_path: Optional[Path],
-        is_py_file: bool = False,
     ) -> None:
         self.ctx = ctx
         self.module_name = module_name
         self.is_init = is_init
         self.file_path = file_path
-        self.is_py_file = is_py_file
+
+    @property
+    def is_py_file(self) -> bool:
+        return self.file_path.suffix == ".py"
 
     def visit_Module(self, node: ast.Module) -> list[NameInfo]:
         return [info for child in node.body for info in self.visit(child)]
@@ -305,15 +301,9 @@ def visit_If(self, node: ast.If) -> Iterable[NameInfo]:
                 yield from self.visit(stmt)
 
     def _visit_condition(self, expr: ast.expr) -> Optional[bool]:
-        visitor = LiteralEvalVisitor(self.ctx, self.file_path)
-        try:
-            value = visitor.visit(expr)
-        except InvalidStub:
-            if not self.is_py_file:
-                raise
-            return None
-        else:
-            return bool(value)
+        return evaluate_expression_truthiness(
+            expr, ctx=self.ctx, file_path=self.file_path
+        )
 
     def visit_Try(self, node: ast.Try) -> Iterable[NameInfo]:
         # try-except sometimes gets used with conditional imports. We assume
@@ -430,17 +420,41 @@ def generic_visit(self, node: ast.AST) -> Iterable[NameInfo]:
         raise InvalidStub(f"Cannot handle node {ast.dump(node)}", self.file_path)
 
 
-class LiteralEvalVisitor(ast.NodeVisitor):
-    """Visitor to evaluate the truthiness of a ``test`` expression in an ``ast.Compare`` node.
+def evaluate_expression_truthiness(
+    expr: ast.expr, *, ctx: SearchContext, file_path: Path
+) -> Optional[bool]:
+    """Attempt to statically evaluate the truthiness of the expression represented by ``expr``.
+
+    This is useful for evaluating conditions that are used for branches in stubs, such as
+    ``if sys.platform == "linux": ...`` or ``if sys.version_info >= (3, 8): ...``. It is usually
+    desirable for a type checker only to consider one of these branches as reachable code for a
+    given configuration of the type checker.
 
-    ``LiteralEvalVisitor(ctx, path).visit(node)`` will return ``True`` if ``node`` is an
-    expression that can be statically determined to always be ``True``, ``False`` if it can
-    be statically determined to always be ``False``, and ``None`` if its truthiness cannot
-    be determined statically. For example, if passed an AST node representing the expression
-    ``sys.platform == "linux"``, it will return ``True`` if ``ctx.platform`` is equal to
-    ``"linux"``, otherwise ``False``.
+    Details:
+    * If the truthiness can be statically determined to always be ``True``, it returns ``True``.
+    * If the truthiness can be statically determined to always be ``False``, it returns ``False``.
+    * If the truthiness cannot be statically determined:
+      * If ``file_path`` has a ``.pyi`` extension, ``InvalidStub`` is raised
+      * If ``file_path`` has a any other extension, however, it returns ``None``, since it is
+        expected that non-stub Python source files may contain dynamic expressions in ``if`` tests
+        that cannot be evaluated statically.
+
+    For example, if passed an AST node representing the expression ``sys.platform == "linux"``,
+    it will return ``True`` if ``ctx.platform`` is equal to ``"linux"``, otherwise ``False``.
     """
 
+    visitor = _LiteralEvalVisitor(ctx, file_path)
+    try:
+        value = visitor.visit(expr)
+    except InvalidStub:
+        if file_path.suffix == ".pyi":
+            raise
+        return None
+    else:
+        return bool(value)
+
+
+class _LiteralEvalVisitor(ast.NodeVisitor):
     def __init__(self, ctx: SearchContext, file_path: Optional[Path]) -> None:
         self.ctx = ctx
         self.file_path = file_path
