refactor(autojac): Revamp OrderedSet (#375)

* Change OrderedSet to compose with OrderedDict instead of inheritting from it
* Replace string type hint by future annotations import

Author: ValerianRey

src/torchjd/_autojac/_transform/_ordered_set.py

@@ -1,32 +1,44 @@
+from __future__ import annotations
+
 from collections import OrderedDict
-from collections.abc import Hashable, Iterable, MutableSet
+from collections.abc import Hashable, Iterable, Iterator, MutableSet
 from typing import TypeVar
 
-_KeyType = TypeVar("_KeyType", bound=Hashable)
+_T = TypeVar("_T", bound=Hashable)
 
 
-class OrderedSet(OrderedDict[_KeyType, None], MutableSet[_KeyType]):
+class OrderedSet(MutableSet[_T]):
     """Ordered collection of distinct elements."""
 
-    def __init__(self, elements: Iterable[_KeyType]):
-        super().__init__([(element, None) for element in elements])
+    def __init__(self, elements: Iterable[_T]):
+        super().__init__()
+        self.ordered_dict = OrderedDict[_T, None]([(element, None) for element in elements])
 
-    def difference_update(self, elements: set[_KeyType]) -> None:
+    def difference_update(self, elements: set[_T]) -> None:
         """Removes all specified elements from the OrderedSet."""
 
         for element in elements:
             self.discard(element)
 
-    def add(self, element: _KeyType) -> None:
+    def add(self, element: _T) -> None:
         """Adds the specified element to the OrderedSet."""
 
-        self[element] = None
+        self.ordered_dict[element] = None
 
-    def __add__(self, other: "OrderedSet[_KeyType]") -> "OrderedSet[_KeyType]":
+    def __add__(self, other: OrderedSet[_T]) -> OrderedSet[_T]:
         """Creates a new OrderedSet with the elements of self followed by the elements of other."""
 
         return OrderedSet([*self, *other])
 
-    def discard(self, value: _KeyType) -> None:
+    def discard(self, value: _T) -> None:
         if value in self:
-            del self[value]
+            del self.ordered_dict[value]
+
+    def __iter__(self) -> Iterator[_T]:
+        return self.ordered_dict.__iter__()
+
+    def __len__(self) -> int:
+        return len(self.ordered_dict)
+
+    def __contains__(self, element: object) -> bool:
+        return element in self.ordered_dict