Skip to content

Commit e21ddfd

Browse files
chore: add generic typing to grid generator (#37)
Signed-off-by: jaapschoutenalliander <[email protected]>
1 parent 899b2c1 commit e21ddfd

File tree

1 file changed

+6
-4
lines changed

1 file changed

+6
-4
lines changed

src/power_grid_model_ds/_core/data_source/generator/grid_generators.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
"""Generators for the grid"""
66

7-
from typing import Type
7+
from typing import Generic, Type, TypeVar
88

99
import numpy as np
1010

@@ -18,13 +18,15 @@
1818

1919
# pylint: disable=too-few-public-methods,too-many-arguments,too-many-positional-arguments
2020

21+
T = TypeVar("T", bound=Grid)
2122

22-
class RadialGridGenerator:
23+
24+
class RadialGridGenerator(Generic[T]):
2325
"""Generates a random but structurally correct radial grid with the given specifications"""
2426

2527
def __init__(
2628
self,
27-
grid_class: Type[Grid],
29+
grid_class: Type[T],
2830
nr_nodes: int = 100,
2931
nr_sources: int = 2,
3032
nr_nops: int = 10,
@@ -36,7 +38,7 @@ def __init__(
3638
self.nr_sources = nr_sources
3739
self.nr_nops = nr_nops
3840

39-
def run(self, seed=None, create_10_3_kv_net: bool = False):
41+
def run(self, seed=None, create_10_3_kv_net: bool = False) -> T:
4042
"""Run the generator to create a random radial grid.
4143
4244
if a seed is provided, this will be used to set rng.

0 commit comments

Comments
 (0)