diff --git a/matpowercaseframes/core.py b/matpowercaseframes/core.py index c75569b..fa71eb1 100644 --- a/matpowercaseframes/core.py +++ b/matpowercaseframes/core.py @@ -30,6 +30,7 @@ def __init__( allow_any_keys=False, update_index=True, columns_templates=None, + reset_index=False, ): """ Load data and initialize the CaseFrames class. @@ -45,6 +46,19 @@ def __init__( load_case_engine (object, optional): External engine used to call MATPOWER `loadcase` (e.g. Octave). Defaults to None. If None, parse data using matpowercaseframes.reader.parse_file. + prefix (str, optional): + Prefix for each attribute when reading from Excel or CSV directory. + Defaults to an empty string. + suffix (str, optional): + Suffix for each attribute when reading from Excel or CSV directory. + Defaults to an empty string. + allow_any_keys (bool, optional): + Whether to allow any keys beyond the predefined ATTRIBUTES. Defaults to + False. + columns_templates (dict, optional): + Custom column templates for DataFrames. Defaults to None. + reset_index (bool, optional): + Whether to reset indices to 0-based numbering. Defaults to False. Raises: TypeError: If the input data format is unsupported. @@ -66,6 +80,8 @@ def __init__( ) if update_index and self._attributes: self._update_index(allow_any_keys=allow_any_keys) + if reset_index: + self.reset_index() def _read_data( self, diff --git a/tests/__init__.py b/tests/__init__.py index e69de29..8d53ad0 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -0,0 +1,18 @@ +import pandas as pd +from pandas.testing import assert_frame_equal, assert_index_equal + + +def assert_cf_equal(cf1, cf2): + for attribute in cf1.attributes: + df1 = getattr(cf1, attribute) + df2 = getattr(cf2, attribute) + if isinstance(df1, pd.DataFrame): + assert_frame_equal(df1, df2) + elif isinstance(df1, pd.Index): + assert_index_equal(df1, df2) + else: + try: + assert df1 == df2 + except ValueError as e: + print(df1, df2) + raise ValueError(e) diff --git a/tests/test_core.py b/tests/test_core.py index 8fc2116..45af838 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -7,6 +7,8 @@ from matpowercaseframes import CaseFrames from matpowercaseframes.idx import BUS_I, BUS_TYPE +from .__init__ import assert_cf_equal + """ pytest -n auto -rA --lf -c pyproject.toml --cov-report term-missing --cov=matpowercaseframes tests/ """ @@ -282,3 +284,13 @@ def test_reset_index_and_infer_numpy_case9(): # gen buses must also reference 0..n-1 assert cf.gen["GEN_BUS"].between(0, len(cf.bus) - 1).all() + + # test reset_index as argument + cf_reset = CaseFrames(CASE_PATH_CASE9, reset_index=True) + cf_reset.infer_numpy() + + assert cf_reset.branch["F_BUS"].between(0, len(cf_reset.bus) - 1).all() + assert cf_reset.branch["T_BUS"].between(0, len(cf_reset.bus) - 1).all() + assert cf_reset.gen["GEN_BUS"].between(0, len(cf_reset.bus) - 1).all() + + assert_cf_equal(cf, cf_reset) diff --git a/tests/test_read_matpower_cases.py b/tests/test_read_matpower_cases.py index defb8b8..8f11a60 100644 --- a/tests/test_read_matpower_cases.py +++ b/tests/test_read_matpower_cases.py @@ -2,31 +2,16 @@ import pandas as pd from matpower import path_matpower, start_instance -from pandas.testing import assert_frame_equal, assert_index_equal from matpowercaseframes import CaseFrames +from .__init__ import assert_cf_equal + """ pytest -n auto -rA --cov-report term --cov=matpowercaseframes tests/ """ -def assert_cf_equal(cf1, cf2): - for attribute in cf1.attributes: - df1 = getattr(cf1, attribute) - df2 = getattr(cf2, attribute) - if isinstance(df1, pd.DataFrame): - assert_frame_equal(df1, df2) - elif isinstance(df1, pd.Index): - assert_index_equal(df1, df2) - else: - try: - assert df1 == df2 - except ValueError as e: - print(df1, df2) - raise ValueError(e) - - def test_case9(): CASE_NAME = "case9.m" cf = CaseFrames(CASE_NAME)