perf: Port frozendict to rust

Author: tobni

src/python/pants/backend/go/util_rules/third_party_pkg.py

@@ -115,7 +115,7 @@ def debug_hint(self) -> str:
 
 
 @dataclass(frozen=True)
-class AllThirdPartyPackages(FrozenDict[str, ThirdPartyPkgAnalysis]):
+class AllThirdPartyPackages:
     """All the packages downloaded from a go.mod, along with a digest of the downloaded files.
 
     The digest has files in the format `gopath/pkg/mod`, which is what `GoSdkProcess` sets `GOPATH`

src/python/pants/backend/url_handlers/s3/register.py

@@ -189,7 +189,7 @@ async def download_from_s3(
         NativeDownloadFile(
             url=virtual_hosted_url,
             expected_digest=request.expected_digest,
-            auth_headers=http_request.headers,
+            auth_headers=dict(http_request.headers),
             retry_delay_duration=global_options.file_downloads_retry_delay,
             max_attempts=global_options.file_downloads_max_attempts,
         )

src/python/pants/engine/internals/native_engine.pyi

@@ -6,7 +6,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable, Iterable, Mapping, Sequence
+from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
 from datetime import datetime
 from io import RawIOBase
 from pathlib import Path
@@ -49,6 +49,38 @@ from pants.engine.process import (
 class PyFailure:
     def get_error(self) -> Exception | None: ...
 
+K = TypeVar("K")
+V = TypeVar("V")
+
+class FrozenDict(Mapping[K, V]):
+    """A wrapper around a normal `dict` that removes all methods to mutate the instance and that
+    implements __hash__.
+
+    This should be used instead of normal dicts when working with the engine because normal dicts
+    are not safe to use.
+    """
+
+    @overload
+    def __new__(cls, __items: Iterable[tuple[K, V]], **kwargs: V) -> Self: ...
+    @overload
+    def __new__(cls, __other: Mapping[K, V], **kwargs: V) -> Self: ...
+    @overload
+    def __new__(cls, **kwargs: V) -> Self: ...
+    @classmethod
+    def deep_freeze(cls, data: Mapping[K, V]) -> Self: ...
+    @staticmethod
+    def frozen(to_freeze: Mapping[K, V]) -> FrozenDict[K, V]: ...
+    def __getitem__(self, k: K) -> V: ...
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[K]: ...
+    def __reversed__(self) -> Iterator[K]: ...
+    def __eq__(self, other: Any) -> Any: ...
+    def __lt__(self, other: Any) -> bool: ...
+    def __or__(self, other: Any) -> FrozenDict[K, V]: ...
+    def __ror__(self, other: Any) -> FrozenDict[K, V]: ...
+    def __hash__(self) -> int: ...
+    def __repr__(self) -> str: ...
+
 # ------------------------------------------------------------------------------
 # Address
 # ------------------------------------------------------------------------------

src/python/pants/engine/target.py

@@ -1214,7 +1214,7 @@ def require_unparametrized_overrides(self) -> dict[str, Mapping[str, Any]]:
 class GeneratedTargets(FrozenDict[Address, Target]):
     """A mapping of the address of generated targets to the targets themselves."""
 
-    def __init__(self, generator: Target, generated_targets: Iterable[Target]) -> None:
+    def __new__(cls, generator: Target, generated_targets: Iterable[Target]) -> GeneratedTargets:
         expected_spec_path = generator.address.spec_path
         expected_tgt_name = generator.address.target_name
         mapping = {}
@@ -1242,7 +1242,7 @@ def __init__(self, generator: Target, generated_targets: Iterable[Target]) -> No
                     "Consider using `request.generator.address.create_generated()`."
                 )
             mapping[tgt.address] = tgt
-        super().__init__(mapping)
+        return super().__new__(cls, mapping)
 
 
 @rule(polymorphic=True)

src/python/pants/util/frozendict.py

@@ -3,168 +3,60 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable, Iterable, Iterator, Mapping
-from typing import Any, Self, TypeVar, cast, overload
+from collections.abc import Callable, ItemsView, Iterable, Mapping, ValuesView
+from typing import TypeVar, cast, overload
 
+from pants.engine.internals.native_engine import FrozenDict as FrozenDict
 from pants.util.memo import memoized_method
-from pants.util.strutil import softwrap
 
 K = TypeVar("K")
 V = TypeVar("V")
-
-
-class FrozenDict(Mapping[K, V]):
-    """A wrapper around a normal `dict` that removes all methods to mutate the instance and that
-    implements __hash__.
-
-    This should be used instead of normal dicts when working with the engine because normal dicts
-    are not safe to use.
-    """
-
-    @overload
-    def __init__(self, __items: Iterable[tuple[K, V]], **kwargs: V) -> None: ...
-
-    @overload
-    def __init__(self, __other: Mapping[K, V], **kwargs: V) -> None: ...
-
-    @overload
-    def __init__(self, **kwargs: V) -> None: ...
-
-    def __init__(self, *item: Mapping[K, V] | Iterable[tuple[K, V]], **kwargs: V) -> None:
-        """Creates a `FrozenDict` with arguments accepted by `dict` that also must be hashable."""
-        if len(item) > 1:
-            raise ValueError(
-                f"{type(self).__name__} was called with {len(item)} positional arguments but it expects one."
-            )
-
-        # NB: Keep the variable name `_data` in sync with `externs/mod.rs`.
-        self._data = dict(item[0]) if item else dict()
-        self._data.update(**kwargs)  # type: ignore[call-overload]
-
-        # NB: We eagerly compute the hash to validate that the values are hashable and to avoid
-        # performing the calculation multiple times. This can be revisited if it's found to be a
-        # performance bottleneck.
-        self._hash = self._calculate_hash()
-
-    @classmethod
-    def deep_freeze(cls, data: Mapping[K, V]) -> Self:
-        """Convert mutable values to their frozen counter parts.
-
-        Sets and lists are turned into tuples and dicts into FrozenDicts.
-        """
-
-        def _freeze(obj):
-            if isinstance(obj, dict):
-                return cls.deep_freeze(obj)
-            if isinstance(obj, (list, set)):
-                return tuple(map(_freeze, obj))
-            return obj
-
-        return cls({k: _freeze(v) for k, v in data.items()})
-
-    @staticmethod
-    def frozen(to_freeze: Mapping[K, V]) -> FrozenDict[K, V]:
-        """Returns a `FrozenDict` containing the keys and values of `to_freeze`.
-
-        If `to_freeze` is already a `FrozenDict`, returns the same object.
-        """
-
-        return to_freeze if isinstance(to_freeze, FrozenDict) else FrozenDict(to_freeze)
-
-    def __getitem__(self, k: K) -> V:
-        return self._data[k]
-
-    def __len__(self) -> int:
-        return len(self._data)
-
-    def __iter__(self) -> Iterator[K]:
-        return iter(self._data)
-
-    def __reversed__(self) -> Iterator[K]:
-        return reversed(tuple(self._data))
-
-    def __eq__(self, other: Any) -> Any:
-        # defer to dict's __eq__
-        return self._data == other
-
-    def __lt__(self, other: Any) -> bool:
-        if not isinstance(other, FrozenDict):
-            return NotImplemented
-        # If sorting each of these on every __lt__ call ends up being a problem we could consider
-        # optimising this, by, for instance, sorting on construction.
-        return sorted(self._data.items()) < sorted(other._data.items())
-
-    def __or__(self, other: Any) -> FrozenDict[K, V]:
-        if isinstance(other, FrozenDict):
-            other = other._data
-        elif not isinstance(other, Mapping):
-            return NotImplemented
-        return FrozenDict(self._data | other)
-
-    def __ror__(self, other: Any) -> FrozenDict[K, V]:
-        if isinstance(other, FrozenDict):
-            other = other._data
-        elif not isinstance(other, Mapping):
-            return NotImplemented
-        return FrozenDict(other | self._data)
-
-    def _calculate_hash(self) -> int:
-        try:
-            h = 0
-            for pair in self._data.items():
-                # xor is commutative, i.e. we get the same hash no matter the order of items. This
-                # "relies" on "hash" of the individual elements being unpredictable enough that such
-                # a naive aggregation is okay. In addition, the Python hash isn't / shouldn't be
-                # used for cryptographically sensitive purposes.
-                h ^= hash(pair)
-            return h
-        except TypeError as e:
-            raise TypeError(
-                softwrap(
-                    f"""
-                    Even though you are using a `{type(self).__name__}`, the underlying values are
-                    not hashable. Please use hashable (and preferably immutable) types for the
-                    underlying values, e.g. use tuples instead of lists and use FrozenOrderedSet
-                    instead of set().
-
-                    Original error message: {e}
-
-                    Value: {self}
-                    """
-                )
-            )
-
-    def __hash__(self) -> int:
-        return self._hash
-
-    def __repr__(self) -> str:
-        return f"{type(self).__name__}({self._data!r})"
+T = TypeVar("T")
 
 
 class LazyFrozenDict(FrozenDict[K, V]):
     """A lazy version of `FrozenDict` where the values are not loaded until referenced."""
 
     @overload
-    def __init__(
-        self, __items: Iterable[tuple[K, Callable[[], V]]], **kwargs: Callable[[], V]
-    ) -> None: ...
+    def __new__(
+        cls, __items: Iterable[tuple[K, Callable[[], V]]], **kwargs: Callable[[], V]
+    ) -> LazyFrozenDict[K, V]: ...
 
     @overload
-    def __init__(self, __other: Mapping[K, Callable[[], V]], **kwargs: Callable[[], V]) -> None: ...
+    def __new__(
+        cls, __other: Mapping[K, Callable[[], V]], **kwargs: Callable[[], V]
+    ) -> LazyFrozenDict[K, V]: ...
 
     @overload
-    def __init__(self, **kwargs: Callable[[], V]) -> None: ...
+    def __new__(cls, **kwargs: Callable[[], V]) -> LazyFrozenDict[K, V]: ...
 
-    def __init__(
-        self,
+    def __new__(
+        cls,
         *item: Mapping[K, Callable[[], V]] | Iterable[tuple[K, Callable[[], V]]],
         **kwargs: Callable[[], V],
-    ) -> None:
-        super().__init__(*item, **kwargs)  # type: ignore[arg-type]
+    ) -> LazyFrozenDict[K, V]:
+        return super().__new__(cls, *item, **kwargs)  # type: ignore[arg-type]
 
     def __getitem__(self, k: K) -> V:
         return self._get_value(k)
 
+    @overload
+    def get(self, key: K, /) -> V | None: ...
+    @overload
+    def get(self, key: K, /, default: T = ...) -> V | T: ...
+
+    def get(self, key: K, /, default: T | None = None) -> V | T | None:
+        try:
+            return self._get_value(key)
+        except KeyError:
+            return default
+
     @memoized_method
     def _get_value(self, k: K) -> V:
-        return cast("Callable[[], V]", self._data[k])()
+        return cast("Callable[[], V]", super().__getitem__(k))()
+
+    def items(self) -> ItemsView[K, V]:
+        return ItemsView(self)
+
+    def values(self) -> ValuesView[V]:
+        return ValuesView(self)