Gather split vertices directly (#16)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vsec"
-version = "0.1.2"
+version = "0.1.3"
 description = "Data structure when some edges in a graph are planar points in essence"
 authors = ["edxu96 <edxu96@outlook.com>"]
 

tests/test_conversion.py

@@ -2,7 +2,7 @@
 import networkx as nx
 from pandas.core.frame import DataFrame
 import pytest as pt
-from vsec.conversion import contract, split
+from vsec.conversion import contract, split, split_vertex_df
 from vsec.geometry import GeoGraph
 
 # There are only two edges left, but sequences might be different.
@@ -43,11 +43,31 @@ def test_split(case_grid: GeoGraph, vertices_grid: DataFrame):
         vertices_grid: vertices in ``case_grid`` and their **type**
             attributes.
     """
+    # Only split 10-0.4 kV transformers.
     vertices_split = vertices_grid.index[vertices_grid["type"] == "STAT1004"]
+
     graph, vertex_df = split(case_grid, vertices_split, NAMING, ATTR, IS_FIRST)
 
     assert isinstance(nx.to_pandas_edgelist(graph), DataFrame)
     assert nx.is_connected(graph)
     assert isinstance(vertex_df, DataFrame)
     assert vertex_df.index.name == "original"
     assert set(vertex_df.columns) == COLUMNS
+    assert vertex_df.shape == (34, 2)
+
+
+@pt.mark.usefixtures("vertices_grid")
+def test_split_vertex_df(vertices_grid: DataFrame):
+    """Check if can gather split vertices and resulted new vertices.
+
+    Args:
+        vertices_grid: vertices in ``case_grid`` and their **type**
+            attributes.
+    """
+    # Only split 10-0.4 kV transformers.
+    vertices_split = vertices_grid.index[vertices_grid["type"] == "STAT1004"]
+
+    res = split_vertex_df(vertices_split, NAMING)
+    assert set(res.columns) == COLUMNS
+    assert res.shape == (34, 2)
+    assert res.index.name == "original"

vsec/conversion.py

@@ -103,3 +103,28 @@ def split(
     vertex_df.index.name = "original"
 
     return graph, vertex_df
+
+
+def split_vertex_df(
+    vertices: Set[str], naming: Callable[[str], Tuple[str, str]],
+) -> DataFrame:
+    """Gather split vertices and resulted new vertices.
+
+    Warning:
+        This should be used when vertex splitting is known to be
+        correct, because there is no validation.
+
+    Args:
+        vertices: vertices ought to be modelled as edges.
+        naming: how two resulted vertices should be named.
+
+    Returns:
+        Correspondence between split vertices & resulted new vertices.
+    """
+    vertex_dict = {}
+    for vertex in vertices:
+        vertex_dict[vertex] = naming(vertex)
+
+    res = pd.DataFrame.from_dict(vertex_dict, orient="index", columns=COLUMNS,)
+    res.index.name = "original"
+    return res