better for loop

sdv-dev · Jun 18, 2024 · c55dd06 · c55dd06
1 parent df28758
commit c55dd06
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/sdv/multi_table/hma.py b/sdv/multi_table/hma.py
@@ -273,11 +273,10 @@ def preprocess(self, data):
 
     def _set_extended_columns_distributions(self, synthesizer, table_name, valid_columns):
         numerical_distributions = {}
-        if table_name in self._parent_extended_columns:
-            for extended_column in self._parent_extended_columns[table_name]:
-                if extended_column in valid_columns:
-                    numerical_distributions[extended_column] = DEFAULT_EXTENDED_COLUMNS_DISTRIBUTION
-            synthesizer._set_numerical_distributions(numerical_distributions)
+        for extended_column in self._parent_extended_columns.get(table_name, []):
+            if extended_column in valid_columns:
+                numerical_distributions[extended_column] = DEFAULT_EXTENDED_COLUMNS_DISTRIBUTION
+        synthesizer._set_numerical_distributions(numerical_distributions)
 
     def _get_extension(self, child_name, child_table, foreign_key, progress_bar_desc):
         """Generate the extension columns for this child table.