Pydantic m.dict() => m.model_dump(); add type hints to conftest.py

devtools/materialize_asset.py

@@ -38,7 +38,7 @@ def main(asset_id):
                 config={
                     "resources": {
                         "dataset_settings": {
-                            "config": etl_fast_settings.dict(),
+                            "config": etl_fast_settings.model_dump(),
                         },
                     },
                 },

src/pudl/cli/etl.py

@@ -26,7 +26,7 @@
 )
 
 import pudl
-from pudl.settings import EtlSettings
+from pudl.settings import EpaCemsSettings, EtlSettings
 from pudl.workspace.setup import PudlPaths
 
 logger = pudl.logging_helpers.get_logger(__name__)
@@ -116,15 +116,15 @@ def main():
 
     etl_settings = EtlSettings.from_yaml(args.settings_file)
 
-    dataset_settings_config = etl_settings.datasets.dict()
+    dataset_settings_config = etl_settings.datasets.model_dump()
     process_epacems = True
     if etl_settings.datasets.epacems is None:
         process_epacems = False
         # Dagster config expects values for the epacems settings even though
         # the CEMS assets will not be executed. Fill in the config dictionary
         # with default cems values. Replace this workaround once dagster pydantic
         # config classes are available.
-        dataset_settings_config["epacems"] = pudl.settings.EpaCemsSettings().dict()
+        dataset_settings_config["epacems"] = EpaCemsSettings().model_dump()
 
     pudl_etl_reconstructable_job = build_reconstructable_job(
         "pudl.cli.etl",

src/pudl/etl/__init__.py

@@ -116,7 +116,7 @@ def load_dataset_settings_from_file(setting_filename: str) -> dict:
     dataset_settings = EtlSettings.from_yaml(
         importlib.resources.files("pudl.package_data.settings")
         / f"{setting_filename}.yml"
-    ).datasets.dict()
+    ).datasets.model_dump()
 
     return dataset_settings
 

src/pudl/ferc_to_sqlite/__init__.py

@@ -51,7 +51,7 @@ def ferc_to_sqlite_xbrl_only():
     config={
         "resources": {
             "ferc_to_sqlite_settings": {
-                "config": ferc_to_sqlite_fast_settings.dict(),
+                "config": ferc_to_sqlite_fast_settings.model_dump(),
             },
         },
     },

src/pudl/ferc_to_sqlite/cli.py

@@ -145,7 +145,7 @@ def main():  # noqa: C901
         run_config={
             "resources": {
                 "ferc_to_sqlite_settings": {
-                    "config": etl_settings.ferc_to_sqlite_settings.dict()
+                    "config": etl_settings.ferc_to_sqlite_settings.model_dump()
                 },
                 "datastore": {
                     "config": {

src/pudl/metadata/classes.py

@@ -158,7 +158,7 @@ class Base(pydantic.BaseModel):
         >>> m.fields
         ['x']
         >>> m.fields = ['y']
-        >>> m.dict()
+        >>> m.model_dump(by_alias=True)
         {'fields': ['y']}
     """
 

src/pudl/transform/ferc1.py

@@ -821,15 +821,16 @@ class GroupMetricChecks(TransformParams):
     group_metric_tolerances: GroupMetricTolerances = GroupMetricTolerances()
     is_close_tolerance: CalculationIsCloseTolerance = CalculationIsCloseTolerance()
 
+    # TODO[pydantic] refactor to work with Pydantic v2
     # @model_validator
     # def grouped_tol_ge_ungrouped_tol(cls, values):
     #     """Grouped tolerance should always be greater than or equal to ungrouped."""
     #     group_metric_tolerances = values["group_metric_tolerances"]
     #     groups_to_check = values["groups_to_check"]
     #     for group in groups_to_check:
-    #         metric_tolerances = group_metric_tolerances.dict().get(group)
+    #         metric_tolerances = group_metric_tolerances.model_dump().get(group)
     #         for metric_name, tolerance in metric_tolerances.items():
-    #             ungrouped_tolerance = group_metric_tolerances.dict()["ungrouped"].get(
+    #             ungrouped_tolerance = group_metric_tolerances.model_dump()["ungrouped"].get(
     #                 metric_name
     #             )
     #             if tolerance < ungrouped_tolerance:
@@ -1175,14 +1176,16 @@ def check_calculation_metrics_by_group(
         for (
             metric_name,
             metric_tolerance,
-        ) in group_metric_checks.group_metric_tolerances.dict()[group_name].items():
+        ) in group_metric_checks.group_metric_tolerances.model_dump()[
+            group_name
+        ].items():
             if metric_name in group_metric_checks.metrics_to_check:
                 # this feels icky. the param name for the metrics are all snake_case while
                 # the metric classes are all TitleCase. So we convert to TitleCase
                 title_case_test = metric_name.title().replace("_", "")
                 group_metric_checker = globals()[title_case_test](
                     by=group_name,
-                    is_close_tolerance=group_metric_checks.is_close_tolerance.dict()[
+                    is_close_tolerance=group_metric_checks.is_close_tolerance.model_dump()[
                         metric_name
                     ],
                     metric_tolerance=metric_tolerance,

test/conftest.py

@@ -5,6 +5,7 @@
 import logging
 import os
 from pathlib import Path
+from typing import Any
 
 import pytest
 import sqlalchemy as sa
@@ -16,13 +17,15 @@
 from pudl.extract.ferc1 import raw_xbrl_metadata_json
 from pudl.ferc_to_sqlite.cli import ferc_to_sqlite_job_factory
 from pudl.io_managers import (
+    PudlSQLiteIOManager,
     ferc1_dbf_sqlite_io_manager,
     ferc1_xbrl_sqlite_io_manager,
     pudl_sqlite_io_manager,
 )
 from pudl.metadata.classes import Package
 from pudl.output.pudltabl import PudlTabl
-from pudl.settings import DatasetsSettings, EtlSettings
+from pudl.settings import DatasetsSettings, EtlSettings, FercToSqliteSettings
+from pudl.workspace.datastore import Datastore
 from pudl.workspace.setup import PudlPaths
 
 logger = logging.getLogger(__name__)
@@ -85,19 +88,19 @@ def test_directory():
 
 
 @pytest.fixture(scope="session", name="live_dbs")
-def live_databases(request):
+def live_databases(request) -> bool:
     """Fixture that tells whether to use existing live FERC1/PUDL DBs)."""
     return request.config.getoption("--live-dbs")
 
 
 @pytest.fixture(scope="session", name="save_unmapped_ids")
-def save_unmapped_ids(request):
+def save_unmapped_ids(request) -> bool:
     """Fixture that tells whether to use existing live FERC1/PUDL DBs)."""
     return request.config.getoption("--save-unmapped-ids")
 
 
 @pytest.fixture(scope="session", name="check_foreign_keys")
-def check_foreign_keys(request):
+def check_foreign_keys(request) -> bool:
     """Fixture that tells whether to use existing live FERC1/PUDL DBs)."""
     return not request.config.getoption("--ignore-foreign-key-constraints")
 
@@ -116,7 +119,7 @@ def etl_parameters(request, test_dir) -> EtlSettings:
 
 
 @pytest.fixture(scope="session", name="ferc_to_sqlite_settings")
-def ferc_to_sqlite_parameters(etl_settings):
+def ferc_to_sqlite_parameters(etl_settings: EtlSettings) -> FercToSqliteSettings:
     """Read ferc_to_sqlite parameters out of test settings dictionary."""
     return etl_settings.ferc_to_sqlite_settings
 
@@ -128,7 +131,7 @@ def pudl_etl_parameters(etl_settings: EtlSettings) -> DatasetsSettings:
 
 
 @pytest.fixture(scope="session", params=["AS"], ids=["ferc1_annual"])
-def pudl_out_ferc1(live_dbs, pudl_engine, request):
+def pudl_out_ferc1(live_dbs: bool, pudl_engine: sa.Engine, request) -> PudlTabl:
     """Define parameterized PudlTabl output object fixture for FERC 1 tests."""
     if not live_dbs:
         pytest.skip("Validation tests only work with a live PUDL DB.")
@@ -140,7 +143,7 @@ def pudl_out_ferc1(live_dbs, pudl_engine, request):
     params=[None, "AS", "MS"],
     ids=["eia_raw", "eia_annual", "eia_monthly"],
 )
-def pudl_out_eia(live_dbs, pudl_engine, request):
+def pudl_out_eia(live_dbs: bool, pudl_engine: sa.Engine, request) -> PudlTabl:
     """Define parameterized PudlTabl output object fixture for EIA tests."""
     if not live_dbs:
         pytest.skip("Validation tests only work with a live PUDL DB.")
@@ -154,9 +157,12 @@ def pudl_out_eia(live_dbs, pudl_engine, request):
 
 
 @pytest.fixture(scope="session", name="fast_out_annual")
-def fast_out_annual(pudl_engine, pudl_datastore_fixture):
+def fast_out_annual(
+    pudl_engine: sa.Engine,
+    pudl_datastore_fixture: Datastore,
+) -> PudlTabl:
     """A PUDL output object for use in CI."""
-    return pudl.output.pudltabl.PudlTabl(
+    return PudlTabl(
         pudl_engine,
         freq="AS",
         fill_fuel_cost=True,
@@ -166,15 +172,17 @@ def fast_out_annual(pudl_engine, pudl_datastore_fixture):
 
 
 @pytest.fixture(scope="session")
-def pudl_out_orig(live_dbs, pudl_engine):
+def pudl_out_orig(live_dbs: bool, pudl_engine: sa.Engine) -> PudlTabl:
     """Create an unaggregated PUDL output object for checking raw data."""
     if not live_dbs:
         pytest.skip("Validation tests only work with a live PUDL DB.")
     return PudlTabl(pudl_engine=pudl_engine)
 
 
 @pytest.fixture(scope="session")
-def ferc_to_sqlite_dbf_only(live_dbs, pudl_datastore_config, etl_settings):
+def ferc_to_sqlite_dbf_only(
+    live_dbs: bool, pudl_datastore_config, etl_settings: EtlSettings
+):
     """Create raw FERC 1 SQLite DBs, but only based on DBF sources."""
     if not live_dbs:
         ferc_to_sqlite_job_factory(
@@ -183,7 +191,7 @@ def ferc_to_sqlite_dbf_only(live_dbs, pudl_datastore_config, etl_settings):
             run_config={
                 "resources": {
                     "ferc_to_sqlite_settings": {
-                        "config": etl_settings.ferc_to_sqlite_settings.dict()
+                        "config": etl_settings.ferc_to_sqlite_settings.model_dump()
                     },
                     "datastore": {
                         "config": pudl_datastore_config,
@@ -194,7 +202,9 @@ def ferc_to_sqlite_dbf_only(live_dbs, pudl_datastore_config, etl_settings):
 
 
 @pytest.fixture(scope="session")
-def ferc_to_sqlite_xbrl_only(live_dbs, pudl_datastore_config, etl_settings):
+def ferc_to_sqlite_xbrl_only(
+    live_dbs: bool, pudl_datastore_config, etl_settings: EtlSettings
+):
     """Create raw FERC 1 SQLite DBs, but only based on XBRL sources."""
     if not live_dbs:
         ferc_to_sqlite_job_factory(
@@ -203,7 +213,7 @@ def ferc_to_sqlite_xbrl_only(live_dbs, pudl_datastore_config, etl_settings):
             run_config={
                 "resources": {
                     "ferc_to_sqlite_settings": {
-                        "config": etl_settings.ferc_to_sqlite_settings.dict()
+                        "config": etl_settings.ferc_to_sqlite_settings.model_dump()
                     },
                     "datastore": {
                         "config": pudl_datastore_config,
@@ -214,7 +224,7 @@ def ferc_to_sqlite_xbrl_only(live_dbs, pudl_datastore_config, etl_settings):
 
 
 @pytest.fixture(scope="session")
-def ferc_to_sqlite(live_dbs, pudl_datastore_config, etl_settings):
+def ferc_to_sqlite(live_dbs, pudl_datastore_config, etl_settings: EtlSettings):
     """Create raw FERC 1 SQLite DBs.
 
     If we are using the test database, we initialize it from scratch first. If we're
@@ -223,14 +233,14 @@ def ferc_to_sqlite(live_dbs, pudl_datastore_config, etl_settings):
     """
     if not live_dbs:
         logger.info(
-            f"ferc_to_sqlite_settings: {etl_settings.ferc_to_sqlite_settings.dict()}"
+            f"ferc_to_sqlite_settings: {etl_settings.ferc_to_sqlite_settings.model_dump()}"
         )
         logger.info(f"ferc_to_sqlite PUDL_OUTPUT: {os.getenv('PUDL_OUTPUT')}")
         ferc_to_sqlite_job_factory()().execute_in_process(
             run_config={
                 "resources": {
                     "ferc_to_sqlite_settings": {
-                        "config": etl_settings.ferc_to_sqlite_settings.dict()
+                        "config": etl_settings.ferc_to_sqlite_settings.model_dump()
                     },
                     "datastore": {
                         "config": pudl_datastore_config,
@@ -241,7 +251,7 @@ def ferc_to_sqlite(live_dbs, pudl_datastore_config, etl_settings):
 
 
 @pytest.fixture(scope="session", name="ferc1_engine_dbf")
-def ferc1_dbf_sql_engine(ferc_to_sqlite_dbf_only):
+def ferc1_dbf_sql_engine(ferc_to_sqlite_dbf_only: FercToSqliteSettings) -> sa.Engine:
     """Grab a connection to the FERC Form 1 DB clone."""
     context = build_init_resource_context(
         resources={"dataset_settings": dataset_settings_config}
@@ -250,7 +260,9 @@ def ferc1_dbf_sql_engine(ferc_to_sqlite_dbf_only):
 
 
 @pytest.fixture(scope="session", name="ferc1_engine_xbrl")
-def ferc1_xbrl_sql_engine(ferc_to_sqlite_xbrl_only, dataset_settings_config):
+def ferc1_xbrl_sql_engine(
+    ferc_to_sqlite_xbrl_only: FercToSqliteSettings, dataset_settings_config
+) -> sa.Engine:
     """Grab a connection to the FERC Form 1 DB clone."""
     context = build_init_resource_context(
         resources={"dataset_settings": dataset_settings_config}
@@ -259,7 +271,7 @@ def ferc1_xbrl_sql_engine(ferc_to_sqlite_xbrl_only, dataset_settings_config):
 
 
 @pytest.fixture(scope="session", name="ferc1_xbrl_taxonomy_metadata")
-def ferc1_xbrl_taxonomy_metadata(ferc1_engine_xbrl):
+def ferc1_xbrl_taxonomy_metadata(ferc1_engine_xbrl: sa.Engine):
     """Read the FERC 1 XBRL taxonomy metadata from JSON."""
     result = materialize_to_memory([raw_xbrl_metadata_json])
     assert result.success
@@ -269,14 +281,14 @@ def ferc1_xbrl_taxonomy_metadata(ferc1_engine_xbrl):
 
 @pytest.fixture(scope="session")
 def pudl_sql_io_manager(
-    ferc1_engine_dbf,  # Implicit dependency
-    ferc1_engine_xbrl,  # Implicit dependency
-    live_dbs,
+    ferc1_engine_dbf: sa.Engine,  # Implicit dependency
+    ferc1_engine_xbrl: sa.Engine,  # Implicit dependency
+    live_dbs: bool,
     pudl_datastore_config,
     dataset_settings_config,
-    check_foreign_keys,
+    check_foreign_keys: bool,
     request,
-):
+) -> PudlSQLiteIOManager:
     """Grab a connection to the PUDL IO manager.
 
     If we are using the test database, we initialize the PUDL DB from scratch. If we're
@@ -309,7 +321,7 @@ def pudl_sql_io_manager(
 
 
 @pytest.fixture(scope="session")
-def pudl_engine(pudl_sql_io_manager):
+def pudl_engine(pudl_sql_io_manager: PudlSQLiteIOManager) -> sa.Engine:
     """Get PUDL SQL engine from io manager."""
     return pudl_sql_io_manager.engine
 
@@ -345,13 +357,13 @@ def configure_paths_for_tests(tmp_path_factory, request):
 
 
 @pytest.fixture(scope="session")
-def dataset_settings_config(request, etl_settings):
+def dataset_settings_config(request, etl_settings: EtlSettings):
     """Create dagster dataset_settings resource."""
-    return etl_settings.datasets.dict()
+    return etl_settings.datasets.model_dump()
 
 
 @pytest.fixture(scope="session")  # noqa: C901
-def pudl_datastore_config(request):
+def pudl_datastore_config(request) -> dict[str, Any]:
     """Produce a :class:pudl.workspace.datastore.Datastore."""
     gcs_cache_path = request.config.getoption("--gcs-cache-path")
     return {
@@ -361,13 +373,13 @@ def pudl_datastore_config(request):
 
 
 @pytest.fixture(scope="session")
-def pudl_datastore_fixture(pudl_datastore_config):
+def pudl_datastore_fixture(pudl_datastore_config: dict[str, Any]) -> Datastore:
     """Create pudl Datastore resource."""
     init_context = build_init_resource_context(config=pudl_datastore_config)
     return resources.datastore(init_context)
 
 
-def skip_table_if_null_freq_table(table_name, freq):
+def skip_table_if_null_freq_table(table_name: str, freq: str | None):
     """Check."""
     if table_name in AS_MS_ONLY_FREQ_TABLES and freq is None:
         pytest.skip(