Adopt generic read_csv function

Author: cc-a

src/muse/new_input/readers.py

@@ -29,7 +29,9 @@ class Demand(TableBase):
 
     year: Mapped[int] = mapped_column(primary_key=True, autoincrement=False)
     commodity: Mapped[Commodities] = mapped_column(
-        ForeignKey("commodities.name"), primary_key=True
+        ForeignKey("commodities.name"),
+        primary_key=True,
+        info=dict(header="commodity_name"),
     )
     region: Mapped[Regions] = mapped_column(
         ForeignKey("regions.name"), primary_key=True
@@ -45,29 +47,29 @@ def read_inputs(data_dir):
     con = engine.raw_connection().driver_connection
 
     with open(data_dir / "regions.csv") as f:
-        regions = read_regions_csv(f, con)  # noqa: F841
+        regions = read_csv(f, Regions, con)  # noqa: F841
 
     with open(data_dir / "commodities.csv") as f:
-        commodities = read_commodities_csv(f, con)
+        commodities = read_csv(f, Commodities, con)
 
     with open(data_dir / "demand.csv") as f:
-        demand = read_demand_csv(f, con)  # noqa: F841
+        demand = read_csv(f, Demand, con)  # noqa: F841
 
     data = {}
     data["global_commodities"] = calculate_global_commodities(commodities)
     return data
 
 
-def read_regions_csv(buffer_, con):
-    rel = con.read_csv(buffer_, header=True, delimiter=",")  # noqa: F841
-    con.execute("INSERT INTO regions SELECT name FROM rel;")
-    return con.sql("SELECT name from regions").fetchnumpy()
-
+def read_csv(buffer_, table_class, con):
+    table_name = table_class.__tablename__
+    columns = ", ".join(
+        column.info.get("header", column.name)
+        for column in table_class.__table__.columns
+    )
 
-def read_commodities_csv(buffer_, con):
     rel = con.read_csv(buffer_, header=True, delimiter=",")  # noqa: F841
-    con.sql("INSERT INTO commodities SELECT name, type, unit FROM rel;")
-    return con.sql("select name, type, unit from commodities").fetchnumpy()
+    con.execute(f"INSERT INTO {table_name} SELECT {columns} FROM rel")
+    return con.execute(f"SELECT * from {table_name}").fetchnumpy()
 
 
 def calculate_global_commodities(commodities):
@@ -85,9 +87,3 @@ def calculate_global_commodities(commodities):
 
     data = xr.Dataset(data_vars=dict(type=type_array, unit=unit_array))
     return data
-
-
-def read_demand_csv(buffer_, con):
-    rel = con.read_csv(buffer_, header=True, delimiter=",")  # noqa: F841
-    con.sql("INSERT INTO demand SELECT year, commodity_name, region, demand FROM rel;")
-    return con.sql("SELECT * from demand").fetchnumpy()

tests/test_readers.py

@@ -341,38 +341,38 @@ def con():
 
 @fixture
 def populate_regions(default_new_input, con):
-    from muse.new_input.readers import read_regions_csv
+    from muse.new_input.readers import Regions, read_csv
 
     with open(default_new_input / "regions.csv") as f:
-        return read_regions_csv(f, con)
+        return read_csv(f, Regions, con)
 
 
 @fixture
 def populate_commodities(default_new_input, con):
-    from muse.new_input.readers import read_commodities_csv
+    from muse.new_input.readers import Commodities, read_csv
 
     with open(default_new_input / "commodities.csv") as f:
-        return read_commodities_csv(f, con)
+        return read_csv(f, Commodities, con)
 
 
 @fixture
 def populate_demand(default_new_input, con, populate_regions, populate_commodities):
-    from muse.new_input.readers import read_demand_csv
+    from muse.new_input.readers import Demand, read_csv
 
     with open(default_new_input / "demand.csv") as f:
-        return read_demand_csv(f, con)
+        return read_csv(f, Demand, con)
 
 
 def test_read_regions(populate_regions):
     assert populate_regions["name"] == np.array(["R1"])
 
 
 def test_read_regions_primary_key_constraint(default_new_input, con):
-    from muse.new_input.readers import read_regions_csv
+    from muse.new_input.readers import Regions, read_csv
 
     csv = StringIO("name\nR1\nR1\n")
     with raises(duckdb.ConstraintException, match=".*duplicate key.*"):
-        read_regions_csv(csv, con)
+        read_csv(csv, Regions, con)
 
 
 def test_read_new_commodities(populate_commodities):
@@ -398,19 +398,19 @@ def test_calculate_global_commodities(populate_commodities):
 
 
 def test_read_new_commodities_primary_key_constraint(default_new_input, con):
-    from muse.new_input.readers import read_commodities_csv
+    from muse.new_input.readers import Commodities, read_csv
 
     csv = StringIO("name,type,unit\nfoo,energy,bar\nfoo,energy,bar\n")
     with raises(duckdb.ConstraintException, match=".*duplicate key.*"):
-        read_commodities_csv(csv, con)
+        read_csv(csv, Commodities, con)
 
 
 def test_read_new_commodities_type_constraint(default_new_input, con):
-    from muse.new_input.readers import read_commodities_csv
+    from muse.new_input.readers import Commodities, read_csv
 
     csv = StringIO("name,type,unit\nfoo,invalid,bar\n")
     with raises(duckdb.ConstraintException):
-        read_commodities_csv(csv, con)
+        read_csv(csv, Commodities, con)
 
 
 def test_new_read_demand_csv(populate_demand):
@@ -424,31 +424,31 @@ def test_new_read_demand_csv(populate_demand):
 def test_new_read_demand_csv_commodity_constraint(
     default_new_input, con, populate_commodities, populate_regions
 ):
-    from muse.new_input.readers import read_demand_csv
+    from muse.new_input.readers import Demand, read_csv
 
     csv = StringIO("year,commodity_name,region,demand\n2020,invalid,R1,0\n")
     with raises(duckdb.ConstraintException, match=".*foreign key.*"):
-        read_demand_csv(csv, con)
+        read_csv(csv, Demand, con)
 
 
 def test_new_read_demand_csv_region_constraint(
     default_new_input, con, populate_commodities, populate_regions
 ):
-    from muse.new_input.readers import read_demand_csv
+    from muse.new_input.readers import Demand, read_csv
 
     csv = StringIO("year,commodity_name,region,demand\n2020,heat,invalid,0\n")
     with raises(duckdb.ConstraintException, match=".*foreign key.*"):
-        read_demand_csv(csv, con)
+        read_csv(csv, Demand, con)
 
 
 def test_new_read_demand_csv_primary_key_constraint(
     default_new_input, con, populate_commodities, populate_regions
 ):
-    from muse.new_input.readers import read_demand_csv, read_regions_csv
+    from muse.new_input.readers import Demand, Regions, read_csv
 
     # Add another region so we can test varying it as a primary key
     csv = StringIO("name\nR2\n")
-    read_regions_csv(csv, con)
+    read_csv(csv, Regions, con)
 
     # all fine so long as one primary key column differs
     csv = StringIO(
@@ -459,12 +459,12 @@ def test_new_read_demand_csv_primary_key_constraint(
 2020,gas,R2,0
 """
     )
-    read_demand_csv(csv, con)
+    read_csv(csv, Demand, con)
 
     # no good if all primary key columns match a previous entry
     csv = StringIO("year,commodity_name,region,demand\n2020,gas,R1,0")
     with raises(duckdb.ConstraintException, match=".*duplicate key.*"):
-        read_demand_csv(csv, con)
+        read_csv(csv, Demand, con)
 
 
 @mark.xfail