add support for both list[str] and str

Signed-off-by: Thijs Baaijen <[email protected]>

Author: Thijss

src/power_grid_model_ds/_core/model/grids/base.py

@@ -409,13 +409,17 @@ def _from_pickle(cls, pickle_path: Path):
         return grid
 
     @classmethod
-    def from_txt(cls, txt_lines: list[str]):
+    def from_txt(cls, txt_lines: list[str] | str):
         """Build a grid from a list of strings
 
         See the documentation for the expected format of the txt_lines
 
-        Example:
+        Args:
+            txt_lines (list[str] | str): The lines of the txt
+
+        Examples:
             >>> Grid.from_txt(["1 2", "2 3", "3 4 transformer", "4 5", "S1 6"])
+            alternative: Grid.from_txt("1 2\n2 3\n3 4 transformer\n4 5\nS1 6")
         """
         return TextSource(grid_class=cls).load_from_txt(txt_lines)
 

tests/unit/model/grids/test_grid_base.py

@@ -293,6 +293,14 @@ def test_from_txt_lines(self):
         assert 1 == grid.transformer.size
         np.testing.assert_array_equal([14, 10, 11, 12, 13, 15, 16, 17], grid.branches.id)
 
+    def test_from_txt_string(self):
+        txt_string = "S1 2\nS1 3 open\n2 7\n3 5\n3 6 transformer\n5 7\n7 8\n8 9"
+        assert Grid.from_txt(txt_string)
+
+    def test_from_txt_string_with_spaces(self):
+        txt_string = "S1 2     \nS1 3   open\n2    7\n3 5\n   3 6 transformer\n5 7\n7   8\n8 9"
+        assert Grid.from_txt(txt_string)
+
     def test_from_txt_with_branch_ids(self):
         txt_lines = [
             "S1 2 91",
@@ -312,25 +320,15 @@ def test_from_txt_with_branch_ids(self):
         np.testing.assert_array_equal([95, 91, 92, 93, 94, 96, 97, 98], grid.branches.id)
 
     def test_from_txt_with_conflicting_ids(self):
-        txt_lines = [
-            "S1 2",
-            "1 3",
-        ]
+        with pytest.raises(ValueError):
+            Grid.from_txt(["S1 2", "1 3"])
 
+    def test_from_txt_with_invalid_line(self):
         with pytest.raises(ValueError):
-            Grid.from_txt(txt_lines)
+            Grid.from_txt(["S1 2 arg3 arg4"])
 
     def test_from_txt_with_unordered_node_ids(self):
-        txt_lines = [
-            "S1 2",
-            "S1 10",
-            "10 11",
-            "2 5",
-            "5 6",
-            "3 4",
-            "3 7",
-        ]
-        grid = Grid.from_txt(txt_lines)
+        grid = Grid.from_txt(["S1 2", "S1 10", "10 11", "2 5", "5 6", "3 4", "3 7"])
         assert 9 == grid.node.size
 
     def test_from_txt_with_unordered_branch_ids(self):