pr feedback

Author: johnnygreco

src/data_designer/config/columns.py

@@ -163,26 +163,28 @@ class SeedDatasetColumnConfig(SingleColumnConfig):
 }
 
 
-def column_type_is_in_dag(column_type: Union[str, DataDesignerColumnType]) -> bool:
+def column_type_used_in_execution_dag(column_type: Union[str, DataDesignerColumnType]) -> bool:
+    """Return True if the column type is used in the workflow execution DAG."""
     column_type = resolve_string_enum(column_type, DataDesignerColumnType)
-    return column_type in [
+    return column_type in {
         DataDesignerColumnType.EXPRESSION,
         DataDesignerColumnType.LLM_CODE,
         DataDesignerColumnType.LLM_JUDGE,
         DataDesignerColumnType.LLM_STRUCTURED,
         DataDesignerColumnType.LLM_TEXT,
         DataDesignerColumnType.VALIDATION,
-    ]
+    }
 
 
 def column_type_is_llm_generated(column_type: Union[str, DataDesignerColumnType]) -> bool:
+    """Return True if the column type is an LLM-generated column."""
     column_type = resolve_string_enum(column_type, DataDesignerColumnType)
-    return column_type in [
+    return column_type in {
         DataDesignerColumnType.LLM_TEXT,
         DataDesignerColumnType.LLM_CODE,
         DataDesignerColumnType.LLM_STRUCTURED,
         DataDesignerColumnType.LLM_JUDGE,
-    ]
+    }
 
 
 def get_column_config_from_kwargs(name: str, column_type: DataDesignerColumnType, **kwargs) -> ColumnConfigT:
@@ -217,6 +219,7 @@ def get_column_config_from_kwargs(name: str, column_type: DataDesignerColumnType
 
 
 def get_column_display_order() -> list[DataDesignerColumnType]:
+    """Return the preferred display order of the column types."""
     return [
         DataDesignerColumnType.SEED_DATASET,
         DataDesignerColumnType.SAMPLER,

src/data_designer/engine/dataset_builders/utils/dag.py

@@ -5,7 +5,7 @@
 
 import networkx as nx
 
-from data_designer.config.columns import ColumnConfigT, column_type_is_in_dag
+from data_designer.config.columns import ColumnConfigT, column_type_used_in_execution_dag
 from data_designer.engine.dataset_builders.utils.errors import DAGCircularDependencyError
 
 logger = logging.getLogger(__name__)
@@ -14,8 +14,12 @@
 def topologically_sort_column_configs(column_configs: list[ColumnConfigT]) -> list[ColumnConfigT]:
     dag = nx.DiGraph()
 
-    non_dag_column_config_list = [col for col in column_configs if not column_type_is_in_dag(col.column_type)]
-    dag_column_config_dict = {col.name: col for col in column_configs if column_type_is_in_dag(col.column_type)}
+    non_dag_column_config_list = [
+        col for col in column_configs if not column_type_used_in_execution_dag(col.column_type)
+    ]
+    dag_column_config_dict = {
+        col.name: col for col in column_configs if column_type_used_in_execution_dag(col.column_type)
+    }
 
     if len(dag_column_config_dict) == 0:
         return non_dag_column_config_list

tests/config/test_columns.py

@@ -15,8 +15,8 @@
     Score,
     SeedDatasetColumnConfig,
     ValidationColumnConfig,
-    column_type_is_in_dag,
     column_type_is_llm_generated,
+    column_type_used_in_execution_dag,
     get_column_config_from_kwargs,
     get_column_display_order,
 )
@@ -56,14 +56,14 @@ def test_data_designer_column_type_is_llm_generated():
 
 
 def test_data_designer_column_type_is_in_dag():
-    assert column_type_is_in_dag(DataDesignerColumnType.EXPRESSION)
-    assert column_type_is_in_dag(DataDesignerColumnType.LLM_CODE)
-    assert column_type_is_in_dag(DataDesignerColumnType.LLM_JUDGE)
-    assert column_type_is_in_dag(DataDesignerColumnType.LLM_STRUCTURED)
-    assert column_type_is_in_dag(DataDesignerColumnType.LLM_TEXT)
-    assert column_type_is_in_dag(DataDesignerColumnType.VALIDATION)
-    assert not column_type_is_in_dag(DataDesignerColumnType.SAMPLER)
-    assert not column_type_is_in_dag(DataDesignerColumnType.SEED_DATASET)
+    assert column_type_used_in_execution_dag(DataDesignerColumnType.EXPRESSION)
+    assert column_type_used_in_execution_dag(DataDesignerColumnType.LLM_CODE)
+    assert column_type_used_in_execution_dag(DataDesignerColumnType.LLM_JUDGE)
+    assert column_type_used_in_execution_dag(DataDesignerColumnType.LLM_STRUCTURED)
+    assert column_type_used_in_execution_dag(DataDesignerColumnType.LLM_TEXT)
+    assert column_type_used_in_execution_dag(DataDesignerColumnType.VALIDATION)
+    assert not column_type_used_in_execution_dag(DataDesignerColumnType.SAMPLER)
+    assert not column_type_used_in_execution_dag(DataDesignerColumnType.SEED_DATASET)
 
 
 def test_sampler_column_config():