diff --git a/src/dbt_osmosis/core/osmosis.py b/src/dbt_osmosis/core/osmosis.py index c9d3edb..a99768d 100644 --- a/src/dbt_osmosis/core/osmosis.py +++ b/src/dbt_osmosis/core/osmosis.py @@ -28,7 +28,11 @@ MissingOsmosisConfig, ) from dbt_osmosis.core.log_controller import logger -from dbt_osmosis.core.column_level_knowledge_propagator import ColumnLevelKnowledgePropagator +from dbt_osmosis.core.column_level_knowledge_propagator import ( + ColumnLevelKnowledgePropagator, + ColumnLevelKnowledge, + Knowledge, +) from dbt_osmosis.vendored.dbt_core_interface.project import ( ColumnInfo, DbtProject, @@ -958,9 +962,9 @@ def remove_columns_not_in_database( @staticmethod def get_prior_knowledge( - knowledge: Dict[str, Dict[str, Any]], + knowledge: Knowledge, column: str, - ) -> Dict[str, Any]: + ) -> ColumnLevelKnowledge: camel_column = re.sub("_(.)", lambda m: m.group(1).upper(), column) prior_knowledge_candidates = list(filter(lambda k: k, [ knowledge.get(column), @@ -990,7 +994,7 @@ def update_undocumented_columns_with_prior_knowledge( ) -> int: """Update undocumented columns with prior knowledge in node and model simultaneously THIS MUTATES THE NODE AND MODEL OBJECTS so that state is always accurate""" - knowledge = ColumnLevelKnowledgePropagator.get_node_columns_with_inherited_knowledge( + knowledge: Knowledge = ColumnLevelKnowledgePropagator.get_node_columns_with_inherited_knowledge( self.manifest, node, self.placeholders ) @@ -1002,7 +1006,7 @@ def update_undocumented_columns_with_prior_knowledge( changes_committed = 0 for column in undocumented_columns: - prior_knowledge = self.get_prior_knowledge(knowledge, column) + prior_knowledge: ColumnLevelKnowledge = self.get_prior_knowledge(knowledge, column) progenitor = prior_knowledge.pop("progenitor", "Unknown") prior_knowledge = {k: v for k, v in prior_knowledge.items() if k in inheritables} if not prior_knowledge: