From 299ec4e84da586ce47a93de2ca3ae69f4dd8dcad Mon Sep 17 00:00:00 2001 From: Courtney Holcomb Date: Tue, 19 Nov 2024 12:51:43 -0800 Subject: [PATCH] fixup! Use helper function to find matching instance in dataset --- metricflow/dataset/sql_dataset.py | 27 ++++++++++++++------------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/metricflow/dataset/sql_dataset.py b/metricflow/dataset/sql_dataset.py index 1acafddaa..59db7290a 100644 --- a/metricflow/dataset/sql_dataset.py +++ b/metricflow/dataset/sql_dataset.py @@ -122,27 +122,28 @@ def column_association_for_dimension( return column_associations_to_return[0] - def instance_for_time_dimension(self, time_dimension_spec: TimeDimensionSpec) -> TimeDimensionInstance: - """Given the name of the time dimension, return the instance associated with it in the data set.""" + def instances_for_time_dimensions( + self, time_dimension_specs: Sequence[TimeDimensionSpec] + ) -> List[TimeDimensionInstance]: + """Return the instances associated with these specs in the data set.""" matching_instances = 0 - instance_to_return = None + instances_to_return: List[TimeDimensionInstance] = [] for time_dimension_instance in self.instance_set.time_dimension_instances: - if time_dimension_instance.spec == time_dimension_spec: - instance_to_return = time_dimension_instance + if time_dimension_instance.spec in time_dimension_specs: + instances_to_return.append(time_dimension_instance) matching_instances += 1 - if matching_instances > 1: + if matching_instances != len(time_dimension_specs): raise RuntimeError( - f"More than one time dimension instance with spec {time_dimension_spec} in " - f"instance set: {self.instance_set}" + f"Unexpected number of time dimension instances found matching specs.\nSpecs: {time_dimension_specs}\n" + f"Instances: {matching_instances}" ) - if not instance_to_return: - raise RuntimeError( - f"No time dimension instances with spec {time_dimension_spec} in instance set: {self.instance_set}" - ) + return instances_to_return - return instance_to_return + def instance_for_time_dimension(self, time_dimension_spec: TimeDimensionSpec) -> TimeDimensionInstance: + """Given the name of the time dimension, return the instance associated with it in the data set.""" + return self.instances_for_time_dimensions([time_dimension_spec])[0] def column_association_for_time_dimension(self, time_dimension_spec: TimeDimensionSpec) -> ColumnAssociation: """Given the name of the time dimension, return the set of columns associated with it in the data set."""