22
33import collections
44from pathlib import Path
5- from typing import List , Optional , Callable , OrderedDict , Dict , Union , Iterable , cast
5+ from typing import Callable , OrderedDict , Iterable , cast
66
77import networkx
88import numpy as np
9+ import numpy .typing as npt
910import pandas as pd
1011import joblib
1112
@@ -40,8 +41,8 @@ class Mesh2VecBase:
4041 def __init__ (
4142 self ,
4243 distance : int ,
43- hyper_edges : Dict [str , List [str ]],
44- vtx_ids : Optional [ List [ str ]] = None ,
44+ hyper_edges : dict [str , list [str ]],
45+ vtx_ids : list [ str ] | None = None ,
4546 calc_strategy : str = "dfs" ,
4647 ):
4748 # pylint: disable=line-too-long
@@ -82,7 +83,7 @@ def __init__(
8283 check_vtx_ids (vtx_ids , hyper_edges )
8384
8485 self ._distance : int = distance
85- self ._hyper_edges : OrderedDict [str , List [str ]] = collections .OrderedDict (hyper_edges )
86+ self ._hyper_edges : OrderedDict [str , list [str ]] = collections .OrderedDict (hyper_edges )
8687
8788 self ._vtx_idx_to_ids = collections .OrderedDict (enumerate (vtx_ids )) # type: ignore
8889 self ._vtx_ids_to_idx = {vtx_ids : i for i , vtx_ids in enumerate (vtx_ids )}
@@ -203,7 +204,7 @@ def from_file(hg_file: Path, distance: int, calc_strategy: str = "dfs") -> "Mesh
203204 hyper_edges_ids_to_vtx_ids = {line_split [0 ]: line_split [1 :] for line_split in lines_split }
204205 return Mesh2VecBase (distance = distance , hyper_edges = hyper_edges_ids_to_vtx_ids )
205206
206- def get_nbh (self , vtx : str , dist : int ) -> List [str ]:
207+ def get_nbh (self , vtx : str , dist : int ) -> list [str ]:
207208 """
208209 Get a list of neighbors with the exact distance ``dist`` of a given vertex ``vtx``
209210
@@ -228,10 +229,10 @@ def get_nbh(self, vtx: str, dist: int) -> List[str]:
228229 def aggregate_categorical (
229230 self ,
230231 feature : str ,
231- dist : Union [ List [ int ], int ] ,
232- categories : Optional [ Union [ List [ str ], List [int ]]] = None ,
233- default_value : Optional [ Union [ int , str ]] = None ,
234- ) -> Union [ str , List [str ] ]:
232+ dist : list [ int ] | int ,
233+ categories : list [ str ] | list [int ] | None = None ,
234+ default_value : int | str | None = None ,
235+ ) -> str | list [str ]:
235236 """
236237 For categorical features, aggregate the numbers of occurrences of each categorical value.
237238 This results in a new aggregated ``feature`` for each categorical value. If ``feature`` is
@@ -263,6 +264,7 @@ def aggregate_categorical(
263264 check_distance_arg (dist_to_check , self )
264265 check_feature_available (feature , self )
265266 feature_names = []
267+ feature_categories : list [str ] | npt .NDArray [np .str_ ]
266268 if categories is not None :
267269 feature_categories = [str (category ) for category in categories ] + ["NONE" ]
268270 else :
@@ -284,15 +286,15 @@ def aggregate_categorical(
284286 def aggregate (
285287 self ,
286288 feature : str ,
287- dist : Union [ List [ int ], int ] ,
288- aggr : Union [
289- Callable [[np .ndarray ], Union [ float , int , str ]],
290- Callable [[np .ndarray , Union [ float , int , str ]], Union [ float , int , str ]],
291- ] ,
292- aggr_name : Optional [ str ] = None ,
289+ dist : list [ int ] | int ,
290+ aggr : (
291+ Callable [[np .ndarray ], float | int | str ]
292+ | Callable [[np .ndarray , float | int | str ], float | int | str ]
293+ ) ,
294+ aggr_name : str | None = None ,
293295 agg_add_ref : bool = False ,
294296 default_value : float = 0.0 ,
295- ) -> Union [ str , List [str ] ]:
297+ ) -> str | list [str ]:
296298 # pylint: disable=line-too-long,too-many-arguments,too-many-positional-arguments
297299 """
298300 Aggregate features from neighborhoods for each distance in ``dist``
@@ -364,15 +366,14 @@ def _collect_feature_values(
364366 self ,
365367 feature_name : str ,
366368 dist : int ,
367- default_value : Optional [Union [float , int , str ]],
368- aggr : Optional [
369- Union [
370- Callable [[np .ndarray ], Union [float , int , str ]],
371- Callable [[np .ndarray , Union [float , int , str ]], Union [float , int , str ]],
372- ]
373- ] = None ,
374- ref_values : Optional [List [Union [float , int , str ]]] = None ,
375- ) -> List [Union [float , int , str ]]:
369+ default_value : float | int | str | None ,
370+ aggr : (
371+ Callable [[np .ndarray ], float | int | str ]
372+ | Callable [[np .ndarray , float | int | str ], float | int | str ]
373+ | None
374+ ) = None ,
375+ ref_values : list [float | int | str ] | None = None ,
376+ ) -> list [float | int | str ]:
376377 """helper method to collect and aggregate data from all hyper nodes"""
377378 # pylint: disable=too-many-arguments,too-many-positional-arguments,line-too-long
378379 check_distance_arg (dist , self )
@@ -384,11 +385,10 @@ def _collect_feature_values(
384385 if dist == 0 :
385386 return feature
386387
387- # `aggr` and `default_value` are explicitly defined in method's signature, so we just
388- # ignore typing here and let it fail in case of wrong usage.
388+ # Ignore typing here and let it fail in case of wrong usage.
389389 if ref_values is None :
390390 if aggr is not None : # fast: use nan_to_num over the whole array
391- aggr = cast (Callable [[np .ndarray ], Union [ float , int , str ] ], aggr )
391+ aggr = cast (Callable [[np .ndarray ], float | int | str ], aggr )
392392 return np .nan_to_num (
393393 [aggr (feature [neighborhood ]) for neighborhood in self ._neighborhoods [dist ]],
394394 nan = default_value , # type: ignore[arg-type]
@@ -406,7 +406,7 @@ def _collect_feature_values(
406406 raise TypeError ("When `ref_values` is given, `aggr` function to compare is needed." )
407407
408408 # compare to reference value is needed
409- aggr = cast (Callable [[np .ndarray , Union [ float , int , str ]], Union [ float , int , str ] ], aggr )
409+ aggr = cast (Callable [[np .ndarray , float | int | str ], float | int | str ], aggr )
410410 return [
411411 np .nan_to_num (
412412 aggr (feature [neighborhood ], ref_values [i ]),
@@ -419,7 +419,7 @@ def add_features_from_csv(
419419 self ,
420420 csv_file : Path ,
421421 with_header : bool = False ,
422- columns : Optional [ List [ str ]] = None ,
422+ columns : list [ str ] | None = None ,
423423 ) -> None :
424424 """Map the content of a CSV file to the vertices of the hypergraph.
425425
@@ -465,7 +465,7 @@ def add_features_from_dataframe(self, df: pd.DataFrame) -> None:
465465 raise ValueError (f"Feature { new_columns_name } already exists" )
466466 self ._features = self ._features .merge (df , how = "left" , on = "vtx_id" , validate = "1:1" )
467467
468- def to_dataframe (self , vertices : Optional [ Iterable [str ]] = None ) -> pd .DataFrame :
468+ def to_dataframe (self , vertices : Iterable [str ] | None = None ) -> pd .DataFrame :
469469 """
470470 Returns a Pandas dataframe with all the beforehand aggregated feature columns. If
471471 ``vertices`` is not ``None`` and iterable, the dataframe is only generated for vertices
@@ -475,7 +475,7 @@ def to_dataframe(self, vertices: Optional[Iterable[str]] = None) -> pd.DataFrame
475475 return self ._aggregated_features [self ._features in vertices ].copy ()
476476 return self ._aggregated_features .copy ()
477477
478- def to_array (self , vertices : Optional [ Iterable [str ]] = None ) -> np .ndarray :
478+ def to_array (self , vertices : Iterable [str ] | None = None ) -> np .ndarray :
479479 """
480480 Returns a numpy array with all the beforehand aggregated feature columns. If
481481 ``vertices`` is not ``None`` and iterable, the array is only generated for vertices
@@ -487,15 +487,15 @@ def get_max_distance(self) -> int:
487487 """returns the distance value used to generate the hypergraph neighborhood"""
488488 return self ._distance
489489
490- def available_features (self ) -> List [str ]:
490+ def available_features (self ) -> list [str ]:
491491 """returns a list the names of all features"""
492492 return self ._features .drop ("vtx_id" , axis = 1 ).keys ().tolist ()
493493
494- def available_aggregated_features (self ) -> List [str ]:
494+ def available_aggregated_features (self ) -> list [str ]:
495495 """returns a list the names of all aggregated features"""
496496 return self ._aggregated_features .drop ("vtx_id" , axis = 1 ).keys ().tolist ()
497497
498- def vtx_ids (self ) -> List [str ]:
498+ def vtx_ids (self ) -> list [str ]:
499499 """returns a list the ids of all hyper vertices"""
500500 return list (self ._vtx_ids_to_idx .keys ())
501501
0 commit comments