Merge pull request #4324 from tybug/typing

Add more type hints

Author: Zac-HD

hypothesis-python/RELEASE.rst

@@ -0,0 +1,3 @@
+RELEASE_TYPE: patch
+
+Improve our internal type hints.

hypothesis-python/src/hypothesis/extra/ghostwriter.py

@@ -648,7 +648,7 @@ def _imports_for_strategy(strategy):
         for f in strategy.flat_conditions:
             imports |= _imports_for_object(f)
     if isinstance(strategy, FlatMapStrategy):
-        imports |= _imports_for_strategy(strategy.flatmapped_strategy)
+        imports |= _imports_for_strategy(strategy.base)
         imports |= _imports_for_object(strategy.expand)
 
     # recurse through one_of to handle e.g. from_type(Optional[Foo])

hypothesis-python/src/hypothesis/internal/conjecture/data.py

@@ -11,7 +11,7 @@
 import math
 import time
 from collections import defaultdict
-from collections.abc import Iterable, Iterator, Sequence
+from collections.abc import Hashable, Iterable, Iterator, Sequence
 from enum import IntEnum
 from functools import cached_property
 from random import Random
@@ -704,6 +704,7 @@ def __init__(
         self._sampled_from_all_strategies_elements_message: Optional[
             tuple[str, object]
         ] = None
+        self._shared_strategy_draws: dict[Hashable, Any] = {}
 
         self.expected_exception: Optional[BaseException] = None
         self.expected_traceback: Optional[str] = None

hypothesis-python/src/hypothesis/internal/conjecture/engine.py

@@ -653,7 +653,7 @@ def test_function(self, data: ConjectureData) -> None:
 
         self.record_for_health_check(data)
 
-    def on_pareto_evict(self, data: ConjectureData) -> None:
+    def on_pareto_evict(self, data: ConjectureResult) -> None:
         self.settings.database.delete(self.pareto_key, choices_to_bytes(data.choices))
 
     def generate_novel_prefix(self) -> tuple[ChoiceT, ...]:

hypothesis-python/src/hypothesis/internal/conjecture/pareto.py

@@ -8,17 +8,28 @@
 # v. 2.0. If a copy of the MPL was not distributed with this file, You can
 # obtain one at https://mozilla.org/MPL/2.0/.
 
+from collections.abc import Iterator
 from enum import Enum
+from random import Random
+from typing import TYPE_CHECKING, Callable, Optional, Union
 
 from sortedcontainers import SortedList
 
 from hypothesis.internal.conjecture.choice import choices_key
-from hypothesis.internal.conjecture.data import ConjectureData, ConjectureResult, Status
+from hypothesis.internal.conjecture.data import (
+    ConjectureData,
+    ConjectureResult,
+    Status,
+    _Overrun,
+)
 from hypothesis.internal.conjecture.junkdrawer import LazySequenceCopy
 from hypothesis.internal.conjecture.shrinker import sort_key
 
 NO_SCORE = float("-inf")
 
+if TYPE_CHECKING:
+    from hypothesis.internal.conjecture.engine import ConjectureRunner
+
 
 class DominanceRelation(Enum):
     NO_DOMINANCE = 0
@@ -27,7 +38,7 @@ class DominanceRelation(Enum):
     RIGHT_DOMINATES = 3
 
 
-def dominance(left, right):
+def dominance(left: ConjectureResult, right: ConjectureResult) -> DominanceRelation:
     """Returns the dominance relation between ``left`` and ``right``, according
     to the rules that one ConjectureResult dominates another if and only if it
     is better in every way.
@@ -125,21 +136,25 @@ class ParetoFront:
     see how much of a problem this is in practice before we try that.
     """
 
-    def __init__(self, random):
+    def __init__(self, random: Random) -> None:
         self.__random = random
-        self.__eviction_listeners = []
+        self.__eviction_listeners: list[Callable[[ConjectureResult], None]] = []
 
-        self.front = SortedList(key=lambda d: sort_key(d.nodes))
-        self.__pending = None
+        self.front: SortedList[ConjectureResult] = SortedList(
+            key=lambda d: sort_key(d.nodes)
+        )
+        self.__pending: Optional[ConjectureResult] = None
 
-    def add(self, data):
+    def add(self, data: Union[ConjectureData, ConjectureResult, _Overrun]) -> bool:
         """Attempts to add ``data`` to the pareto front. Returns True if
         ``data`` is now in the front, including if data is already in the
         collection, and False otherwise"""
         if data.status < Status.VALID:
             return False
 
+        assert not isinstance(data, _Overrun)
         data = data.as_result()
+        assert not isinstance(data, _Overrun)
 
         if not self.front:
             self.front.add(data)
@@ -166,7 +181,7 @@ def add(self, data):
             # We track which values we are going to remove and remove them all
             # at the end so the shape of the front doesn't change while we're
             # using it.
-            to_remove = []
+            to_remove: list[ConjectureResult] = []
 
             # We now iteratively sample elements from the approximate pareto
             # front to check whether they should be retained. When the set of
@@ -238,26 +253,26 @@ def add(self, data):
         finally:
             self.__pending = None
 
-    def on_evict(self, f):
+    def on_evict(self, f: Callable[[ConjectureResult], None]) -> None:
         """Register a listener function that will be called with data when it
         gets removed from the front because something else dominates it."""
         self.__eviction_listeners.append(f)
 
-    def __contains__(self, data):
+    def __contains__(self, data: object) -> bool:
         return isinstance(data, (ConjectureData, ConjectureResult)) and (
             data.as_result() in self.front
         )
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[ConjectureResult]:
         return iter(self.front)
 
-    def __getitem__(self, i):
+    def __getitem__(self, i: int) -> ConjectureResult:
         return self.front[i]
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.front)
 
-    def __remove(self, data):
+    def __remove(self, data: ConjectureResult) -> None:
         try:
             self.front.remove(data)
         except ValueError:
@@ -277,11 +292,12 @@ class ParetoOptimiser:
     grow more powerful over time.
     """
 
-    def __init__(self, engine):
+    def __init__(self, engine: "ConjectureRunner") -> None:
         self.__engine = engine
-        self.front = self.__engine.pareto_front
+        assert self.__engine.pareto_front is not None
+        self.front: ParetoFront = self.__engine.pareto_front
 
-    def run(self):
+    def run(self) -> None:
         seen = set()
 
         # We iterate backwards through the pareto front, using the shrinker to

hypothesis-python/src/hypothesis/internal/reflection.py

@@ -23,12 +23,13 @@
 import warnings
 from collections.abc import MutableMapping, Sequence
 from functools import partial, wraps
+from inspect import Parameter, Signature
 from io import StringIO
 from keyword import iskeyword
 from random import _inst as global_random_instance
 from tokenize import COMMENT, detect_encoding, generate_tokens, untokenize
 from types import ModuleType
-from typing import Any, Callable, TypeVar
+from typing import Any, Callable, Optional, TypeVar
 from unittest.mock import _patch as PatchType
 from weakref import WeakKeyDictionary
 
@@ -43,7 +44,7 @@
 LAMBDA_SOURCE_CACHE: MutableMapping[Callable, str] = WeakKeyDictionary()
 
 
-def is_mock(obj):
+def is_mock(obj: object) -> bool:
     """Determine if the given argument is a mock type."""
 
     # We want to be able to detect these when dealing with various test
@@ -118,7 +119,7 @@ def function_digest(function: Any) -> bytes:
     return hasher.digest()
 
 
-def check_signature(sig: inspect.Signature) -> None:
+def check_signature(sig: Signature) -> None:
     # Backport from Python 3.11; see https://github.com/python/cpython/pull/92065
     for p in sig.parameters.values():
         if iskeyword(p.name) and p.kind is not p.POSITIONAL_ONLY:
@@ -132,17 +133,19 @@ def check_signature(sig: inspect.Signature) -> None:
 
 def get_signature(
     target: Any, *, follow_wrapped: bool = True, eval_str: bool = False
-) -> inspect.Signature:
+) -> Signature:
     # Special case for use of `@unittest.mock.patch` decorator, mimicking the
     # behaviour of getfullargspec instead of reporting unusable arguments.
     patches = getattr(target, "patchings", None)
     if isinstance(patches, list) and all(isinstance(p, PatchType) for p in patches):
-        P = inspect.Parameter
-        return inspect.Signature(
-            [P("args", P.VAR_POSITIONAL), P("keywargs", P.VAR_KEYWORD)]
+        return Signature(
+            [
+                Parameter("args", Parameter.VAR_POSITIONAL),
+                Parameter("keywargs", Parameter.VAR_KEYWORD),
+            ]
         )
 
-    if isinstance(getattr(target, "__signature__", None), inspect.Signature):
+    if isinstance(getattr(target, "__signature__", None), Signature):
         # This special case covers unusual codegen like Pydantic models
         sig = target.__signature__
         check_signature(sig)
@@ -152,7 +155,7 @@ def get_signature(
             selfy = next(iter(sig.parameters.values()))
             if (
                 selfy.name == "self"
-                and selfy.default is inspect.Parameter.empty
+                and selfy.default is Parameter.empty
                 and selfy.kind.name.startswith("POSITIONAL_")
             ):
                 return sig.replace(
@@ -172,10 +175,10 @@ def get_signature(
     return sig
 
 
-def arg_is_required(param):
-    return param.default is inspect.Parameter.empty and param.kind in (
-        inspect.Parameter.POSITIONAL_OR_KEYWORD,
-        inspect.Parameter.KEYWORD_ONLY,
+def arg_is_required(param: Parameter) -> bool:
+    return param.default is Parameter.empty and param.kind in (
+        Parameter.POSITIONAL_OR_KEYWORD,
+        Parameter.KEYWORD_ONLY,
     )
 
 
@@ -204,7 +207,9 @@ def required_args(target, args=(), kwargs=()):
     }
 
 
-def convert_keyword_arguments(function, args, kwargs):
+def convert_keyword_arguments(
+    function: Any, args: Sequence[object], kwargs: dict[str, object]
+) -> tuple[tuple[object, ...], dict[str, object]]:
     """Returns a pair of a tuple and a dictionary which would be equivalent
     passed as positional and keyword args to the function. Unless function has
     kwonlyargs or **kwargs the dictionary will always be empty.
@@ -238,24 +243,22 @@ def convert_positional_arguments(
     return tuple(new_args), new_kwargs
 
 
-def ast_arguments_matches_signature(args, sig):
-    assert isinstance(args, ast.arguments)
-    assert isinstance(sig, inspect.Signature)
-    expected = []
+def ast_arguments_matches_signature(args: ast.arguments, sig: Signature) -> bool:
+    expected: list[tuple[str, int]] = []
     for node in args.posonlyargs:
-        expected.append((node.arg, inspect.Parameter.POSITIONAL_ONLY))
+        expected.append((node.arg, Parameter.POSITIONAL_ONLY))
     for node in args.args:
-        expected.append((node.arg, inspect.Parameter.POSITIONAL_OR_KEYWORD))
+        expected.append((node.arg, Parameter.POSITIONAL_OR_KEYWORD))
     if args.vararg is not None:
-        expected.append((args.vararg.arg, inspect.Parameter.VAR_POSITIONAL))
+        expected.append((args.vararg.arg, Parameter.VAR_POSITIONAL))
     for node in args.kwonlyargs:
-        expected.append((node.arg, inspect.Parameter.KEYWORD_ONLY))
+        expected.append((node.arg, Parameter.KEYWORD_ONLY))
     if args.kwarg is not None:
-        expected.append((args.kwarg.arg, inspect.Parameter.VAR_KEYWORD))
+        expected.append((args.kwarg.arg, Parameter.VAR_KEYWORD))
     return expected == [(p.name, p.kind) for p in sig.parameters.values()]
 
 
-def is_first_param_referenced_in_function(f):
+def is_first_param_referenced_in_function(f: Any) -> bool:
     """Is the given name referenced within f?"""
     try:
         tree = ast.parse(textwrap.dedent(inspect.getsource(f)))
@@ -301,7 +304,7 @@ def _extract_lambda_source(f):
     # The answer is that we add this at runtime, in new_given_signature(),
     # and we do support strange choices as applying @given() to a lambda.
     sig = inspect.signature(f)
-    assert sig.return_annotation in (inspect.Parameter.empty, None), sig
+    assert sig.return_annotation in (Parameter.empty, None), sig
 
     # Using pytest-xdist on Python 3.13, there's an entry in the linecache for
     # file "<string>", which then returns nonsense to getsource.  Discard it.
@@ -459,7 +462,7 @@ def get_pretty_function_description(f: object) -> str:
     return name
 
 
-def nicerepr(v):
+def nicerepr(v: Any) -> str:
     if inspect.isfunction(v):
         return get_pretty_function_description(v)
     elif isinstance(v, type):
@@ -500,15 +503,15 @@ def repr_call(
     return rep + "(" + ", ".join(bits) + ")"
 
 
-def check_valid_identifier(identifier):
+def check_valid_identifier(identifier: str) -> None:
     if not identifier.isidentifier():
         raise ValueError(f"{identifier!r} is not a valid python identifier")
 
 
-eval_cache: dict = {}
+eval_cache: dict[str, ModuleType] = {}
 
 
-def source_exec_as_module(source):
+def source_exec_as_module(source: str) -> ModuleType:
     try:
         return eval_cache[source]
     except KeyError:
@@ -532,7 +535,9 @@ def {name}{signature}:
 """.lstrip()
 
 
-def get_varargs(sig, kind=inspect.Parameter.VAR_POSITIONAL):
+def get_varargs(
+    sig: Signature, kind: int = Parameter.VAR_POSITIONAL
+) -> Optional[Parameter]:
     for p in sig.parameters.values():
         if p.kind is kind:
             return p
@@ -583,7 +588,7 @@ def accept(f):
         for p in signature.parameters.values():
             if p.kind is p.KEYWORD_ONLY:
                 invocation_parts.append(f"{p.name}={p.name}")
-        varkw = get_varargs(signature, kind=inspect.Parameter.VAR_KEYWORD)
+        varkw = get_varargs(signature, kind=Parameter.VAR_KEYWORD)
         if varkw:
             invocation_parts.append("**" + varkw.name)
 

hypothesis-python/src/hypothesis/strategies/_internal/core.py

@@ -79,6 +79,7 @@
     get_type_hints,
     is_typed_named_tuple,
 )
+from hypothesis.internal.conjecture.data import ConjectureData
 from hypothesis.internal.conjecture.utils import (
     calc_label_from_cls,
     check_sample,
@@ -2199,7 +2200,7 @@ class DataObject:
     # Note that "only exists" here really means "is only exported to users",
     # but we want to treat it as "semi-stable", not document it as "public API".
 
-    def __init__(self, data):
+    def __init__(self, data: ConjectureData) -> None:
         self.count = 0
         self.conjecture_data = data