1212 Pattern ,
1313 Tuple ,
1414 Type ,
15+ TypeVar ,
1516 Union ,
1617)
1718
1819from .object import Object
1920
2021
21- class DictList (list ):
22+ CobraObject = TypeVar ("CobraObject" , bound = Object )
23+
24+
25+ class DictList (List [CobraObject ]):
2226 """
2327 Define a combined dict and list.
2428
@@ -49,12 +53,12 @@ def __init__(self, *args):
4953 self .extend (other )
5054
5155 # noinspection PyShadowingBuiltins
52- def has_id (self , id : Union [Object , str ]) -> bool :
56+ def has_id (self , id : Union [CobraObject , str ]) -> bool :
5357 """Check if id is in DictList."""
5458 return id in self ._dict
5559
5660 # noinspection PyShadowingBuiltins
57- def _check (self , id : Union [Object , str ]) -> None :
61+ def _check (self , id : Union [CobraObject , str ]) -> None :
5862 """Make sure duplicate id's are not added.
5963
6064 This function is called before adding in elements.
@@ -68,15 +72,17 @@ def _generate_index(self) -> None:
6872 self ._dict = {v .id : k for k , v in enumerate (self )}
6973
7074 # noinspection PyShadowingBuiltins
71- def get_by_id (self , id : Union [Object , str ]) -> Object :
75+ def get_by_id (self , id : Union [CobraObject , str ]) -> CobraObject :
7276 """Return the element with a matching id."""
7377 return list .__getitem__ (self , self ._dict [id ])
7478
7579 def list_attr (self , attribute : str ) -> list :
7680 """Return a list of the given attribute for every object."""
7781 return [getattr (i , attribute ) for i in self ]
7882
79- def get_by_any (self , iterable : List [Union [str , Object , int ]]) -> list :
83+ def get_by_any (
84+ self , iterable : List [Union [str , CobraObject , int ]]
85+ ) -> List [CobraObject ]:
8086 """Get a list of members using several different ways of indexing.
8187
8288 Parameters
@@ -92,7 +98,7 @@ def get_by_any(self, iterable: List[Union[str, Object, int]]) -> list:
9298 a list of members
9399 """
94100
95- def get_item (item : Any ) -> Any :
101+ def get_item (item : Any ) -> CobraObject :
96102 if isinstance (item , int ):
97103 return self [item ]
98104 elif isinstance (item , str ):
@@ -110,7 +116,7 @@ def query(
110116 self ,
111117 search_function : Union [str , Pattern , Callable ],
112118 attribute : Union [str , None ] = None ,
113- ) -> "DictList" :
119+ ) -> "DictList[CobraObject] " :
114120 """Query the list.
115121
116122 Parameters
@@ -167,29 +173,29 @@ def select_attribute(x: Optional[Any]) -> Any:
167173 results ._extend_nocheck (matches )
168174 return results
169175
170- def _replace_on_id (self , new_object : Object ) -> None :
176+ def _replace_on_id (self , new_object : CobraObject ) -> None :
171177 """Replace an object by another with the same id."""
172178 the_id = new_object .id
173179 the_index = self ._dict [the_id ]
174180 list .__setitem__ (self , the_index , new_object )
175181
176182 # overriding default list functions with new ones
177- def append (self , entity : Object ) -> None :
183+ def append (self , entity : CobraObject ) -> None :
178184 """Append object to end."""
179185 the_id = entity .id
180186 self ._check (the_id )
181187 self ._dict [the_id ] = len (self )
182188 list .append (self , entity )
183189
184- def union (self , iterable : Iterable [Object ]) -> None :
190+ def union (self , iterable : Iterable [CobraObject ]) -> None :
185191 """Add elements with id's not already in the model."""
186192 _dict = self ._dict
187193 append = self .append
188194 for i in iterable :
189195 if i .id not in _dict :
190196 append (i )
191197
192- def extend (self , iterable : Iterable [Object ]) -> None :
198+ def extend (self , iterable : Iterable [CobraObject ]) -> None :
193199 """Extend list by appending elements from the iterable.
194200
195201 Sometimes during initialization from an older pickle, _dict
@@ -222,7 +228,7 @@ def extend(self, iterable: Iterable[Object]) -> None:
222228 f"Is it present twice?"
223229 )
224230
225- def _extend_nocheck (self , iterable : Iterable [Object ]) -> None :
231+ def _extend_nocheck (self , iterable : Iterable [CobraObject ]) -> None :
226232 """Extend without checking for uniqueness.
227233
228234 This function should only be used internally by DictList when it
@@ -244,7 +250,7 @@ def _extend_nocheck(self, iterable: Iterable[Object]) -> None:
244250 for i , obj in enumerate (islice (self , current_length , None ), current_length ):
245251 _dict [obj .id ] = i
246252
247- def __sub__ (self , other : Iterable [Object ]) -> "DictList" :
253+ def __sub__ (self , other : Iterable [CobraObject ]) -> "DictList[CobraObject] " :
248254 """Remove a value or values, and returns the new DictList.
249255
250256 x.__sub__(y) <==> x - y
@@ -264,7 +270,7 @@ def __sub__(self, other: Iterable[Object]) -> "DictList":
264270 total .remove (item )
265271 return total
266272
267- def __isub__ (self , other : Iterable [Object ]) -> "DictList" :
273+ def __isub__ (self , other : Iterable [CobraObject ]) -> "DictList[CobraObject] " :
268274 """Remove a value or values in place.
269275
270276 x.__sub__(y) <==> x -= y
@@ -278,7 +284,7 @@ def __isub__(self, other: Iterable[Object]) -> "DictList":
278284 self .remove (item )
279285 return self
280286
281- def __add__ (self , other : Iterable [Object ]) -> "DictList" :
287+ def __add__ (self , other : Iterable [CobraObject ]) -> "DictList[CobraObject] " :
282288 """Add item while returning a new DictList.
283289
284290 x.__add__(y) <==> x + y
@@ -294,7 +300,7 @@ def __add__(self, other: Iterable[Object]) -> "DictList":
294300 total .extend (other )
295301 return total
296302
297- def __iadd__ (self , other : Iterable [Object ]) -> "DictList" :
303+ def __iadd__ (self , other : Iterable [CobraObject ]) -> "DictList[CobraObject] " :
298304 """Add item while returning the same DictList.
299305
300306 x.__iadd__(y) <==> x += y
@@ -309,7 +315,7 @@ def __iadd__(self, other: Iterable[Object]) -> "DictList":
309315 self .extend (other )
310316 return self
311317
312- def __reduce__ (self ) -> Tuple [Type ["DictList" ], Tuple , dict , Iterator ]:
318+ def __reduce__ (self ) -> Tuple [Type ["DictList" ], Tuple , dict , Iterator [ CobraObject ] ]:
313319 """Return a reduced version of DictList.
314320
315321 This reduced version details the class, an empty Tuple, a dictionary of the
@@ -336,7 +342,7 @@ def __setstate__(self, state: dict) -> None:
336342 self ._generate_index ()
337343
338344 # noinspection PyShadowingBuiltins
339- def index (self , id : Union [str , Object ], * args ) -> int :
345+ def index (self , id : Union [str , CobraObject ], * args ) -> int :
340346 """Determine the position in the list.
341347
342348 Parameters
@@ -360,7 +366,7 @@ def index(self, id: Union[str, Object], *args) -> int:
360366 except KeyError :
361367 raise ValueError (f"{ str (id )} not found" )
362368
363- def __contains__ (self , entity : Union [str , Object ]) -> bool :
369+ def __contains__ (self , entity : Union [str , CobraObject ]) -> bool :
364370 """Ask if the DictList contain an entity.
365371
366372 DictList.__contains__(entity) <==> entity in DictList
@@ -377,14 +383,14 @@ def __contains__(self, entity: Union[str, Object]) -> bool:
377383 the_id = entity
378384 return the_id in self ._dict
379385
380- def __copy__ (self ) -> "DictList" :
386+ def __copy__ (self ) -> "DictList[CobraObject] " :
381387 """Copy the DictList into a new one."""
382388 the_copy = DictList ()
383389 list .extend (the_copy , self )
384390 the_copy ._dict = self ._dict .copy ()
385391 return the_copy
386392
387- def insert (self , index : int , entity : Object ) -> None :
393+ def insert (self , index : int , entity : CobraObject ) -> None :
388394 """Insert entity before index."""
389395 self ._check (entity .id )
390396 list .insert (self , index , entity )
@@ -395,7 +401,7 @@ def insert(self, index: int, entity: Object) -> None:
395401 _dict [i ] = j + 1
396402 _dict [entity .id ] = index
397403
398- def pop (self , * args ) -> Object :
404+ def pop (self , * args ) -> CobraObject :
399405 """Remove and return item at index (default last)."""
400406 value = list .pop (self , * args )
401407 index = self ._dict .pop (value .id )
@@ -409,11 +415,11 @@ def pop(self, *args) -> Object:
409415 _dict [i ] = j - 1
410416 return value
411417
412- def add (self , x : Object ) -> None :
418+ def add (self , x : CobraObject ) -> None :
413419 """Opposite of `remove`. Mirrors set.add."""
414420 self .extend ([x ])
415421
416- def remove (self , x : Union [str , Object ]) -> None :
422+ def remove (self , x : Union [str , CobraObject ]) -> None :
417423 """.. warning :: Internal use only.
418424
419425 Each item is unique in the list which allows this
@@ -445,8 +451,8 @@ def key(i):
445451 self ._generate_index ()
446452
447453 def __getitem__ (
448- self , i : Union [int , slice , Iterable , Object , "DictList" ]
449- ) -> Union ["DictList" , Object ]:
454+ self , i : Union [int , slice , Iterable , CobraObject , "DictList[CobraObject] " ]
455+ ) -> Union ["DictList[CobraObject] " , CobraObject ]:
450456 """Get item from DictList."""
451457 if isinstance (i , int ):
452458 return list .__getitem__ (self , i )
@@ -465,7 +471,9 @@ def __getitem__(
465471 else :
466472 return list .__getitem__ (self , i )
467473
468- def __setitem__ (self , i : Union [slice , int ], y : Union [list , Object ]) -> None :
474+ def __setitem__ (
475+ self , i : Union [slice , int ], y : Union [List [CobraObject ], CobraObject ]
476+ ) -> None :
469477 """Set an item via index or slice.
470478
471479 Parameters
@@ -507,19 +515,21 @@ def __delitem__(self, index: int) -> None:
507515 if j > index :
508516 _dict [i ] = j - 1
509517
510- def __getslice__ (self , i : int , j : int ) -> "DictList" :
518+ def __getslice__ (self , i : int , j : int ) -> "DictList[CobraObject] " :
511519 """Get a slice from it to j of DictList."""
512520 return self .__getitem__ (slice (i , j ))
513521
514- def __setslice__ (self , i : int , j : int , y : Union [list , Object ]) -> None :
522+ def __setslice__ (
523+ self , i : int , j : int , y : Union [List [CobraObject ], CobraObject ]
524+ ) -> None :
515525 """Set slice, where y is an iterable."""
516526 self .__setitem__ (slice (i , j ), y )
517527
518528 def __delslice__ (self , i : int , j : int ) -> None :
519529 """Remove slice."""
520530 self .__delitem__ (slice (i , j ))
521531
522- def __getattr__ (self , attr : Any ) -> Any :
532+ def __getattr__ (self , attr : Any ) -> CobraObject :
523533 """Get an attribute by id."""
524534 try :
525535 return DictList .get_by_id (self , attr )
0 commit comments