Skip to content

Commit cd606bc

Browse files
refactor(autojac): Make OrderedSet a MutableSet (#352)
* Make OrderedSet inherit from MutableSet rather than Set * Extract the discard method in OrderedSet (which is mandatory as it is an abstract method of the MutableSet class)
1 parent 926be31 commit cd606bc

File tree

1 file changed

+7
-4
lines changed

1 file changed

+7
-4
lines changed

src/torchjd/_autojac/_transform/_ordered_set.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
from collections import OrderedDict
2-
from collections.abc import Hashable, Iterable, Set
2+
from collections.abc import Hashable, Iterable, MutableSet
33
from typing import TypeVar
44

55
_KeyType = TypeVar("_KeyType", bound=Hashable)
66

77

8-
class OrderedSet(OrderedDict[_KeyType, None], Set[_KeyType]):
8+
class OrderedSet(OrderedDict[_KeyType, None], MutableSet[_KeyType]):
99
"""Ordered collection of distinct elements."""
1010

1111
def __init__(self, elements: Iterable[_KeyType]):
@@ -15,8 +15,7 @@ def difference_update(self, elements: set[_KeyType]) -> None:
1515
"""Removes all specified elements from the OrderedSet."""
1616

1717
for element in elements:
18-
if element in self:
19-
del self[element]
18+
self.discard(element)
2019

2120
def add(self, element: _KeyType) -> None:
2221
"""Adds the specified element to the OrderedSet."""
@@ -27,3 +26,7 @@ def __add__(self, other: "OrderedSet[_KeyType]") -> "OrderedSet[_KeyType]":
2726
"""Creates a new OrderedSet with the elements of self followed by the elements of other."""
2827

2928
return OrderedSet([*self, *other])
29+
30+
def discard(self, value: _KeyType) -> None:
31+
if value in self:
32+
del self[value]

0 commit comments

Comments
 (0)