python-constricter 0.2.4__py3-none-any.whl → 0.2.5__py3-none-any.whl

constricter/__init__.py

@@ -25,7 +25,7 @@ from constricter.offences import (
 )
 from constricter.rules.checker import Coverage, annotation_coverage, check_source, check_tree
 
-__version__ = "0.2.4"
+__version__ = "0.2.5"
 __all__ = [
     "COMMENT_TYPED_TARGET",
     "DEFAULT_CHECKS",

constricter/cli/command.py

@@ -10,11 +10,11 @@ import json
 import sys
 import tokenize
 from collections.abc import Callable, Mapping, Sequence
-from concurrent.futures import ProcessPoolExecutor
+from concurrent.futures import Future, ProcessPoolExecutor
 from dataclasses import dataclass, field, replace
 from functools import partial
 from pathlib import Path
-from typing import Final, NamedTuple, TextIO, TypeAlias, cast
+from typing import Final, NamedTuple, Self, TextIO, TypeAlias, cast
 
 from constricter import notebook
 from constricter.cli import baseline, hints
@@ -30,6 +30,7 @@ from constricter.offences import (
     Edit,
     Offence,
 )
+from constricter.rules import parsed
 from constricter.rules.checker import Coverage, annotation_coverage, check_source
 
 EXIT_CLEAN: Final = 0
@@ -203,6 +204,10 @@ class _CoverageRun:
 
 
 _FileRun: TypeAlias = _CheckRun | _BaselineRun | _CoverageRun
+_Check: TypeAlias = Callable[[Path, Outside], _FileRun]  # check (or count, or baseline) one file
+# A worker's answers, one per file of its share: what it read (to index), and what it found.
+_Reading: TypeAlias = Future[list[project.Module | None]]
+_Checking: TypeAlias = Future[list[_FileRun]]
 HINT_ROUNDS: Final = 4  # with `--fix --infer-with`: how many times each file is fixed, at most
 
 
@@ -412,14 +417,135 @@ def _checked_all(
         return list(
             itertools.starmap(check, zip(paths, _outside(modules, paths, hinted), strict=True)),
         ), modules
-    pool: ProcessPoolExecutor
-    with ProcessPoolExecutor(max_workers=options.jobs) as pool:
+    workers: _Workers
+    with _Workers(paths, options.jobs) as workers:
         if modules is None:
-            # The index, as the checks, read one file per task.
-            modules = (
-                project.Index({}, []) if coverage else project.index(paths, partial(pool.map, chunksize=16))
+            modules = project.Index({}, []) if coverage else workers.index()
+        return workers.check(check, _outside(modules, paths, hinted)), modules
+
+
+class _Workers:
+    """Worker processes, each with its own share of the files, which it indexes and then checks.
+
+    The trees a worker parsed to index its files are still there to check (`parsed.keep`), within its
+    part of the budget: sharing the files out anew for the check (a pool's way) would parse them again.
+    """
+
+    def __init__(self, paths: Sequence[Path], jobs: int) -> None:
+        """Share `paths` out among up to `jobs` workers (see `_shares`)."""
+        self.paths: Sequence[Path] = paths
+        self.shares: list[list[int]] = _shares(paths, jobs)
+        self.stack: contextlib.ExitStack = contextlib.ExitStack()
+        self.pools: list[ProcessPoolExecutor] = []
+
+    def __enter__(self) -> Self:
+        """Start the workers, one process each, each with its part of the budget.
+
+        Returns:
+          Them.
+
+        """
+        self.pools = [
+            self.stack.enter_context(
+                ProcessPoolExecutor(max_workers=1, initializer=parsed.budget, initargs=(len(self.shares),)),
             )
-        return list(pool.map(check, paths, _outside(modules, paths, hinted))), modules
+            for _ in self.shares
+        ]
+        return self
+
+    def __exit__(self, *_details: object) -> None:
+        """Stop the workers."""
+        self.stack.close()
+
+    def index(self) -> project.Index:
+        """Index every file, each worker its share.
+
+        Returns:
+          The index, as one process would build it.
+
+        """
+        read: list[_Reading] = [
+            pool.submit(_read_share, [self.paths[at] for at in share])
+            for pool, share in zip(self.pools, self.shares, strict=True)
+        ]
+        found: dict[int, project.Module | None] = {}
+        share: list[int]
+        future: _Reading
+        for share, future in zip(self.shares, read, strict=True):
+            found.update(zip(share, future.result(), strict=True))
+        return project.indexed(found[at] for at in range(len(self.paths)))  # in their order, as one at a time
+
+    def check(self, check: _Check, outside: Sequence[Outside]) -> list[_FileRun]:
+        """Check every file, each worker the share it indexed.
+
+        Returns:
+          What each file found, in their order.
+
+        """
+        checked: list[_Checking] = [
+            pool.submit(_check_share, check, [(self.paths[at], outside[at]) for at in share])
+            for pool, share in zip(self.pools, self.shares, strict=True)
+        ]
+        found: dict[int, _FileRun] = {}
+        share: list[int]
+        future: _Checking
+        for share, future in zip(self.shares, checked, strict=True):
+            found.update(zip(share, future.result(), strict=True))
+        return [found[at] for at in range(len(self.paths))]
+
+
+def _shares(paths: Sequence[Path], workers: int) -> list[list[int]]:
+    """Share the files among the workers, the biggest first to the least loaded: a check's time goes by size.
+
+    Returns:
+      Each worker's files, by their place in `paths` (no more workers than files).
+
+    """
+    loads: list[int] = [0] * min(workers, len(paths))
+    shares: list[list[int]] = [[] for _ in loads]
+    sizes: list[int] = [_size(path) for path in paths]
+    index: int
+    for index in sorted(range(len(paths)), key=lambda at: -sizes[at]):
+        least: int = loads.index(min(loads))
+        shares[least].append(index)
+        loads[least] += sizes[index] or 1
+    return [sorted(share) for share in shares if share]
+
+
+def _size(path: Path) -> int:
+    """Find a file's size, to balance the workers by.
+
+    Returns:
+      It, or 0 if it can't be read (checking it reports that).
+
+    """
+    try:
+        return path.stat().st_size
+    except OSError:
+        return 0
+
+
+def _read_share(paths: Sequence[Path]) -> list[project.Module | None]:
+    """Index a worker's share of the files (see `project.read`).
+
+    Returns:
+      Each file's module, or `None`.
+
+    """
+    return [project.read(path) for path in paths]
+
+
+def _check_share(
+    check: _Check,
+    files: Sequence[tuple[Path, Outside]],
+) -> list[_FileRun]:
+    """Check a worker's share of the files.
+
+    Returns:
+      What each found.
+
+    """
+    return list(itertools.starmap(check, files))
 
 
 def _merged(before: _CheckRun, after: _CheckRun) -> _CheckRun:
@@ -466,7 +592,9 @@ def _outside(
     path: Path
     for path in paths:
         imported: project.Imported = project.imported(modules, path)
-        found.append(Outside(imported.calls, imported.classes, hinted.get(path, ())))
+        found.append(
+            Outside(imported.calls, imported.classes, hinted.get(path, ()), project.type_vars(modules, path)),
+        )
     return found
 
 

constricter/fix/doubts.py

@@ -0,0 +1,288 @@
+# SPDX-License-Identifier: MIT
+"""Where a type checker sees a value otherwise than `--fix` infers it, and what `--fix` does then.
+
+- A copy, attribute or subscript of a union may be narrowed where it's read (`if x is not None:`,
+  `isinstance`), which a type checker sees and `--fix` doesn't; so may a comprehension's elements, by
+  its condition (`[c for c in cs if isinstance(c, Column)]`). Their fixes are guesses.
+- An ALL_CAPS module-level name is a constant to pyright, which keeps its literal's type
+  (`Literal["r"]`) where `str` would widen it: a guess too, where the module passes it to a call
+  or a default (see `passed`), for a parameter that may take only some values.
+- `self`, and a method declared to return `Self` called on `self` or `cls`, are `Self`, not the class
+  (a subclass's `self` isn't its base): written as the module already spells `Self`, or not at all.
+- A generic class the module defines, written bare (`list[Mapper]`, as `self` makes one), is missing
+  its type arguments: not offered.
+"""
+
+import ast
+from collections.abc import Mapping
+from functools import lru_cache
+from types import MappingProxyType
+from typing import Final, NamedTuple, TypeAlias, cast
+
+from constricter.fix.known import ImportPlan, Inference
+from constricter.rules.annotations import node_name
+from constricter.rules.flow import members
+from constricter.rules.syntax import FunctionDef
+from constricter.rules.walked import of_type
+
+# What a copy, an attribute and a subscript of a narrowable union rest on, as guesses.
+_Read: TypeAlias = type[ast.expr]
+_READS: Final[dict[_Read, str]] = {
+    ast.Name: "copy",
+    ast.Attribute: "attribute",
+    ast.Subscript: "subscript",
+}
+# Where a function tests a value, a type checker may narrow it (`isinstance`, `is None`, a `TypeGuard`).
+_TESTS: Final = (ast.If, ast.While, ast.Assert, ast.IfExp, ast.comprehension, ast.Match)
+_SELF: Final = "Self"
+_COMPREHENSION: Final = "comprehension"
+_LITERAL: Final = "literal"
+_COMPREHENSIONS: Final = (ast.ListComp, ast.SetComp, ast.DictComp, ast.GeneratorExp)
+_UNIONS: Final = frozenset({"Optional", "Union"})
+_SELF_ORIGINS: Final = frozenset({"typing.Self", "typing_extensions.Self"})
+_TYPING: Final = frozenset({"typing", "typing_extensions"})
+
+
+class Facts(NamedTuple):
+    """What of a module decides what a type checker sees otherwise than `--fix` infers.
+
+    Each class's methods declared to return a bare `Self` (`self_returns`), its generic classes
+    (`generic_classes`), and the names it passes to a call or a default (`passed`).
+    """
+
+    selfish: Mapping[str, frozenset[str]] = MappingProxyType({})
+    generics: frozenset[str] = frozenset()
+    passed: frozenset[str] = frozenset()
+
+
+class Owner(NamedTuple):
+    """The class a method belongs to, the name its instance (or class) goes by, and its `Self` methods."""
+
+    name: str
+    first: str  # `self`, or a classmethod's first parameter
+    selfish: frozenset[str]  # its methods declared to return a bare `Self` (see `self_returns`)
+
+
+def doubts(value: ast.expr, found: Inference, *, constant: bool, narrowed: bool) -> frozenset[str]:
+    """Find what makes `found`, a certain inference of `value`, a guess (see the module docstring).
+
+    `constant`: whether the value is bound to an ALL_CAPS module-level name; `narrowed`: whether the
+    function tests the value itself somewhere (`isinstance(x, C)`, `x is None`, `is_c(x)`), where a
+    type checker narrows it, whatever its type.
+
+    Returns:
+      The guessing mechanisms (`FIX_KINDS`) it then rests on; none if it stays certain.
+
+    """
+    kind: str | None = _READS.get(type(value))
+    if kind is not None and (narrowed or len(members(found.annotation) or ()) > 1):
+        return frozenset({kind})
+    if isinstance(value, _COMPREHENSIONS) and any(g.ifs for g in value.generators) and _has_union(found):
+        return frozenset({_COMPREHENSION})
+    if constant and _literal(value):
+        return frozenset({_LITERAL})
+    return frozenset()
+
+
+@lru_cache(maxsize=64)  # asked by each assignment in the function's scope
+def tested(function: FunctionDef) -> frozenset[str]:
+    """List what a function tests (see `_TESTS`), where a type checker may narrow it: `x`, `self.x`, `d[k]`.
+
+    Returns:
+      Each name, attribute and subscript in a test, as source text.
+
+    """
+    return frozenset(
+        ast.unparse(read)
+        for node in ast.walk(function)
+        if isinstance(node, _TESTS)
+        for test in _tested_parts(node)
+        for read in ast.walk(test)
+        if isinstance(read, ast.Name | ast.Attribute | ast.Subscript)
+    )
+
+
+def _tested_parts(node: ast.AST) -> list[ast.expr]:
+    """Find what one of `_TESTS` tests: an `if`'s condition, a comprehension's, a `match`'s subject.
+
+    Returns:
+      Them.
+
+    """
+    test: ast.expr
+    ifs: list[ast.expr]
+    match node:
+        case ast.If(test=test) | ast.While(test=test) | ast.Assert(test=test) | ast.IfExp(test=test):
+            return [test]
+        case ast.comprehension(ifs=ifs):
+            return list(ifs)
+        case _:  # a `match` (see `_TESTS`)
+            return [cast("ast.Match", node).subject]
+
+
+def says_self(function: FunctionDef) -> bool:
+    """Check whether a function's signature mentions `Self` (`-> Self`, `other: Self`, `list[Self]`).
+
+    In any other method, mypy takes `self` for its class, not `Self`.
+
+    Returns:
+      Whether it does.
+
+    """
+    args: ast.arguments = function.args
+    annotations: list[ast.expr | None] = [
+        function.returns,
+        *(
+            arg.annotation
+            for arg in (*args.posonlyargs, *args.args, *args.kwonlyargs, args.vararg, args.kwarg)
+            if arg
+        ),
+    ]
+    return any(
+        node_name(node) == _SELF
+        for annotation in annotations
+        if annotation is not None
+        for node in ast.walk(annotation)
+        if isinstance(node, ast.Name | ast.Attribute)
+    )
+
+
+def passed(tree: ast.Module) -> frozenset[str]:
+    """Find the names a module passes to a call (`f(X)`, `f(code=X)`) or gives a parameter as its default.
+
+    Where a constant's `Literal` type matters: a parameter may take only some values
+    (`code: Literal["invalid-type", ...]`), which `str` doesn't say.
+
+    Returns:
+      Them.
+
+    """
+    names: set[str] = set()
+    node: ast.AST
+    for node in of_type(tree, ast.Call, ast.arguments):
+        values: list[ast.expr | None] = (
+            [*node.args, *(keyword.value for keyword in node.keywords)]
+            if isinstance(node, ast.Call)
+            else [*cast("ast.arguments", node).defaults, *cast("ast.arguments", node).kw_defaults]
+        )
+        names.update(value.id for value in values if isinstance(value, ast.Name))
+    return frozenset(names)
+
+
+def is_constant(name: str) -> bool:
+    """Check whether pyright takes a name for a constant: letters all capitals (`MAX`, `_DEFAULT_MODE`).
+
+    Returns:
+      Whether it does.
+
+    """
+    return name.upper() == name and any(char.isalpha() for char in name)
+
+
+def corrected(
+    value: ast.expr,
+    found: Inference,
+    owner: Owner | None,
+    generics: frozenset[str],
+    plan: ImportPlan | None,
+) -> Inference | None:
+    """Write `found` as a type checker sees it: `Self` where the value is one; nothing for a bare generic.
+
+    Returns:
+      The inference to offer, or `None` if there's none to.
+
+    """
+    if owner is not None and _selfish(value, found, owner):
+        spelled: str | None = None if plan is None else spelled_self(plan)
+        return (
+            None if spelled is None else found._replace(annotation=spelled, reason=f"{found.reason}: `Self`")
+        )
+    return None if _bare(found.annotation, generics) else found
+
+
+def _selfish(value: ast.expr, found: Inference, owner: Owner) -> bool:
+    """Check whether `value` is `self`, or a `Self` method called on `self` or `cls`, typed as the class.
+
+    Returns:
+      Whether it is.
+
+    """
+    name: str
+    method: str
+    match value:
+        case ast.Name(id=name) if name == owner.first:
+            return found.annotation == owner.name
+        case ast.Call(func=ast.Attribute(value=ast.Name(id=name), attr=method)) if name == owner.first:
+            return method in owner.selfish and found.annotation == owner.name
+        case _:
+            return False
+
+
+def spelled_self(plan: ImportPlan) -> str | None:
+    """Name `Self` as the module already imports it (`Self`, `typing.Self`, `te.Self`).
+
+    Not an import added for it: `typing.Self` is Python 3.11's, and the module may run on older ones.
+
+    Returns:
+      The name, or `None` if it doesn't import it.
+
+    """
+    bound: str
+    origin: str
+    for bound, origin in plan.bound.items():
+        if origin in _SELF_ORIGINS:
+            return bound
+    for bound, origin in plan.bound.items():
+        if origin in _TYPING:
+            return f"{bound}.Self"
+    return None
+
+
+def _bare(annotation: str, generics: frozenset[str]) -> bool:
+    """Check whether an annotation names one of `generics` without subscripting it.
+
+    Returns:
+      Whether it does.
+
+    """
+    if not generics:
+        return False
+    # `annotation` is always `ast.unparse`'s own output, so it's always valid Python to parse back.
+    tree: ast.expr = ast.parse(annotation, mode="eval").body
+    subscripted: set[int] = {id(node.value) for node in ast.walk(tree) if isinstance(node, ast.Subscript)}
+    return any(
+        isinstance(node, ast.Name) and node.id in generics and id(node) not in subscripted
+        for node in ast.walk(tree)
+    )
+
+
+def _has_union(found: Inference) -> bool:
+    """Check whether an annotation has a union anywhere in it (`list[int | None]`, `Optional[str]`).
+
+    Returns:
+      Whether it has.
+
+    """
+    # `found.annotation` is always `ast.unparse`'s own output, so it's always valid Python to parse back.
+    return any(
+        (isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr))
+        or (isinstance(node, ast.Name) and node.id in _UNIONS)
+        or (isinstance(node, ast.Attribute) and node.attr in _UNIONS)
+        for node in ast.walk(ast.parse(found.annotation, mode="eval"))
+    )
+
+
+def _literal(value: ast.expr) -> bool:
+    """Check whether `value` is a literal whose own type a `Literal` can be (a string, a number, a `bool`).
+
+    Returns:
+      Whether it is.
+
+    """
+    match value:
+        case ast.Constant(value=str() | bytes() | int()):  # `bool` is an `int`
+            return True
+        case ast.UnaryOp(op=ast.USub(), operand=ast.Constant(value=int())):
+            return True
+        case _:
+            return False