Flowfile 0.4.1__py3-none-any.whl → 0.5.3__py3-none-any.whl

flowfile_core/flowfile/schema_callbacks.py

@@ -1,22 +1,17 @@
-
-from typing import List
-
-from polars import datatypes
 import polars as pl
-
 from pl_fuzzy_frame_match.output_column_name_utils import set_name_in_fuzzy_mappings
 from pl_fuzzy_frame_match.pre_process import rename_fuzzy_right_mapping
+from polars import datatypes
 
-from flowfile_core.flowfile.flow_data_engine.subprocess_operations.subprocess_operations import fetch_unique_values
 from flowfile_core.configs.flow_logger import main_logger
 from flowfile_core.flowfile.flow_data_engine.flow_file_column.main import FlowfileColumn, PlType
-from flowfile_core.schemas import transform_schema
-from flowfile_core.schemas import input_schema
+from flowfile_core.flowfile.flow_data_engine.subprocess_operations.subprocess_operations import fetch_unique_values
+from flowfile_core.schemas import input_schema, transform_schema
 
 
-def _ensure_all_columns_have_select(left_cols: List[str],
-                                    right_cols: List[str],
-                                    fuzzy_match_input: transform_schema.FuzzyMatchInput):
+def _ensure_all_columns_have_select(
+    left_cols: list[str], right_cols: list[str], fuzzy_match_input: transform_schema.FuzzyMatchInputManager
+):
     """
     Ensure that all columns in the left and right FlowDataEngines are included in the fuzzy match input's select
      statements.
@@ -32,13 +27,16 @@ def _ensure_all_columns_have_select(left_cols: List[str],
     left_cols_in_select = {c.old_name for c in fuzzy_match_input.left_select.renames}
 
     fuzzy_match_input.left_select.renames.extend(
-        [transform_schema.SelectInput(col) for col in left_cols if col not in left_cols_in_select])
+        [transform_schema.SelectInput(col) for col in left_cols if col not in left_cols_in_select]
+    )
     fuzzy_match_input.right_select.renames.extend(
         [transform_schema.SelectInput(col) for col in right_cols if col not in right_cols_in_select]
     )
 
 
-def _order_join_inputs_based_on_col_order(col_order: List[str], join_inputs: transform_schema.JoinInputs) -> None:
+def _order_join_inputs_based_on_col_order(
+    col_order: list[str], join_inputs: transform_schema.JoinInputsManager
+) -> None:
     """
     Ensure that the select columns in the fuzzy match input match the order of the incoming columns.
     This function modifies the join_inputs object in-place.
@@ -46,27 +44,35 @@ def _order_join_inputs_based_on_col_order(col_order: List[str], join_inputs: tra
     Returns:
         None
     """
-    select_map = {select.new_name: select for select in join_inputs.renames}
+    select_map = {select.old_name: select for select in join_inputs.renames}
     ordered_renames = [select_map[col] for col in col_order if col in select_map]
-    join_inputs.renames = ordered_renames
-
-
-def calculate_fuzzy_match_schema(fm_input: transform_schema.FuzzyMatchInput,
-                                 left_schema: List[FlowfileColumn],
-                                 right_schema: List[FlowfileColumn]):
-    _ensure_all_columns_have_select(left_cols=[col.column_name for col in left_schema],
-                                    right_cols=[col.column_name for col in right_schema],
-                                    fuzzy_match_input=fm_input)
-    _order_join_inputs_based_on_col_order(col_order=[col.column_name for col in left_schema],
-                                          join_inputs=fm_input.left_select)
-    _order_join_inputs_based_on_col_order(col_order=[col.column_name for col in right_schema],
-                                          join_inputs=fm_input.right_select)
+    join_inputs.select_inputs.renames = ordered_renames
+
+
+def calculate_fuzzy_match_schema(
+    fm_input: transform_schema.FuzzyMatchInputManager,
+    left_schema: list[FlowfileColumn],
+    right_schema: list[FlowfileColumn],
+):
+    _ensure_all_columns_have_select(
+        left_cols=[col.column_name for col in left_schema],
+        right_cols=[col.column_name for col in right_schema],
+        fuzzy_match_input=fm_input,
+    )
+
+    _order_join_inputs_based_on_col_order(
+        col_order=[col.column_name for col in left_schema], join_inputs=fm_input.left_select
+    )
+    _order_join_inputs_based_on_col_order(
+        col_order=[col.column_name for col in right_schema], join_inputs=fm_input.right_select
+    )
     for column in fm_input.left_select.renames:
         if column.join_key:
             column.keep = True
     for column in fm_input.right_select.renames:
         if column.join_key:
             column.keep = True
+
     left_schema_dict, right_schema_dict = ({ls.name: ls for ls in left_schema}, {rs.name: rs for rs in right_schema})
     fm_input.auto_rename()
     right_renames = {column.old_name: column.new_name for column in fm_input.right_select.renames}
@@ -75,20 +81,27 @@ def calculate_fuzzy_match_schema(fm_input: transform_schema.FuzzyMatchInput,
     for column in fm_input.left_select.renames:
         column_schema = left_schema_dict.get(column.old_name)
         if column_schema and (column.keep or column.join_key):
-            output_schema.append(FlowfileColumn.from_input(column.new_name, column_schema.data_type,
-                                                           example_values=column_schema.example_values))
+            output_schema.append(
+                FlowfileColumn.from_input(
+                    column.new_name, column_schema.data_type, example_values=column_schema.example_values
+                )
+            )
     for column in fm_input.right_select.renames:
         column_schema = right_schema_dict.get(column.old_name)
         if column_schema and (column.keep or column.join_key):
-            output_schema.append(FlowfileColumn.from_input(column.new_name, column_schema.data_type,
-                                                           example_values=column_schema.example_values))
+            output_schema.append(
+                FlowfileColumn.from_input(
+                    column.new_name, column_schema.data_type, example_values=column_schema.example_values
+                )
+            )
     set_name_in_fuzzy_mappings(new_join_mapping)
-    output_schema.extend([FlowfileColumn.from_input(fuzzy_mapping.output_column_name, 'Float64')
-                          for fuzzy_mapping in new_join_mapping])
+    output_schema.extend(
+        [FlowfileColumn.from_input(fuzzy_mapping.output_column_name, "Float64") for fuzzy_mapping in new_join_mapping]
+    )
     return output_schema
 
 
-def get_schema_of_column(node_input_schema: List[FlowfileColumn], col_name: str) -> FlowfileColumn|None:
+def get_schema_of_column(node_input_schema: list[FlowfileColumn], col_name: str) -> FlowfileColumn | None:
     for s in node_input_schema:
         if s.name == col_name:
             return s
@@ -96,44 +109,54 @@ def get_schema_of_column(node_input_schema: List[FlowfileColumn], col_name: str)
 
 class InvalidSetup(ValueError):
     """Error raised when pivot column has too many unique values."""
+
     pass
 
 
 def get_output_data_type_pivot(schema: FlowfileColumn, agg_type: str) -> datatypes:
-    if agg_type in ('count', 'n_unique'):
+    if agg_type in ("count", "n_unique"):
         output_type = datatypes.Float64  # count is always float
-    elif schema.generic_datatype() == 'numeric':
+    elif schema.generic_datatype() == "numeric":
         output_type = datatypes.Float64
-    elif schema.generic_datatype() == 'string':
+    elif schema.generic_datatype() == "string":
         output_type = datatypes.Utf8
-    elif schema.generic_datatype() == 'date':
+    elif schema.generic_datatype() == "date":
         output_type = datatypes.Datetime
     else:
         output_type = datatypes.Utf8
     return output_type
 
 
-def pre_calculate_pivot_schema(node_input_schema: List[FlowfileColumn],
-                               pivot_input: transform_schema.PivotInput,
-                               output_fields: List[input_schema.MinimalFieldInfo] = None,
-                               input_lf: pl.LazyFrame = None) -> List[FlowfileColumn]:
-    index_columns_schema = [get_schema_of_column(node_input_schema, index_col) for index_col in
-                            pivot_input.index_columns]
+def pre_calculate_pivot_schema(
+    node_input_schema: list[FlowfileColumn],
+    pivot_input: transform_schema.PivotInput,
+    output_fields: list[input_schema.MinimalFieldInfo] = None,
+    input_lf: pl.LazyFrame = None,
+) -> list[FlowfileColumn]:
+    index_columns_schema = [
+        get_schema_of_column(node_input_schema, index_col) for index_col in pivot_input.index_columns
+    ]
     val_column_schema = get_schema_of_column(node_input_schema, pivot_input.value_col)
     if output_fields is not None and len(output_fields) > 0:
-        return index_columns_schema+[FlowfileColumn(PlType(column_name=output_field.name,
-                                                           pl_datatype=output_field.data_type)) for output_field in
-                                     output_fields]
+        return index_columns_schema + [
+            FlowfileColumn(PlType(column_name=output_field.name, pl_datatype=output_field.data_type))
+            for output_field in output_fields
+        ]
 
     else:
         max_unique_vals = 200
-        unique_vals = fetch_unique_values(input_lf.select(pivot_input.pivot_column)
-                                          .unique()
-                                          .sort(pivot_input.pivot_column)
-                                          .limit(max_unique_vals).cast(pl.String))
+        unique_vals = fetch_unique_values(
+            input_lf.select(pivot_input.pivot_column)
+            .unique()
+            .sort(pivot_input.pivot_column)
+            .limit(max_unique_vals)
+            .cast(pl.String)
+        )
         if len(unique_vals) >= max_unique_vals:
-            main_logger.warning('Pivot column has too many unique values. Please consider using a different column.'
-                                f' Max unique values: {max_unique_vals}')
+            main_logger.warning(
+                "Pivot column has too many unique values. Please consider using a different column."
+                f" Max unique values: {max_unique_vals}"
+            )
         pl_output_fields = []
         for val in unique_vals:
             if len(pivot_input.aggregations) == 1:
@@ -142,5 +165,5 @@ def pre_calculate_pivot_schema(node_input_schema: List[FlowfileColumn],
             else:
                 for agg in pivot_input.aggregations:
                     output_type = get_output_data_type_pivot(val_column_schema, agg)
-                    pl_output_fields.append(PlType(column_name=f'{val}_{agg}', pl_datatype=output_type))
+                    pl_output_fields.append(PlType(column_name=f"{val}_{agg}", pl_datatype=output_type))
         return index_columns_schema + [FlowfileColumn(pl_output_field) for pl_output_field in pl_output_fields]

flowfile_core/flowfile/setting_generator/__init__.py

@@ -1,2 +1 @@
-
 from flowfile_core.flowfile.setting_generator.settings import setting_generator, setting_updator

flowfile_core/flowfile/setting_generator/setting_generator.py

@@ -1,5 +1,6 @@
+from collections.abc import Callable
+
 from flowfile_core.configs import logger
-from typing import Callable
 
 
 class SettingGenerator:
@@ -13,10 +14,10 @@ class SettingGenerator:
         setattr(self, f.__name__, f)
 
     def get_setting_generator(self, node_type: str) -> Callable:
-        logger.info('getting setting generator for ' + node_type)
+        logger.info("getting setting generator for " + node_type)
 
         if node_type in self.setting_generator_set:
-            logger.info('setting generator found')
+            logger.info("setting generator found")
             return getattr(self, node_type)
         else:
             return lambda x: x
@@ -33,9 +34,9 @@ class SettingUpdator:
         setattr(self, f.__name__, f)
 
     def get_setting_updator(self, node_type: str) -> Callable:
-        logger.info('getting setting updator for ' + node_type)
+        logger.info("getting setting updator for " + node_type)
         if node_type in self.setting_updator_set:
-            logger.info('setting updator found')
+            logger.info("setting updator found")
             return getattr(self, node_type)
         else:
             return lambda x: x

flowfile_core/flowfile/setting_generator/settings.py

@@ -1,10 +1,11 @@
+from collections.abc import Callable, Iterable
+from functools import wraps
+
+from pl_fuzzy_frame_match.models import FuzzyMapping
 
+from flowfile_core.flowfile.setting_generator.setting_generator import SettingGenerator, SettingUpdator
 from flowfile_core.schemas import input_schema, transform_schema
-from typing import Callable, Iterable
-from functools import wraps
 from flowfile_core.schemas.output_model import NodeData
-from flowfile_core.flowfile.setting_generator.setting_generator import SettingGenerator, SettingUpdator
-from pl_fuzzy_frame_match.models import FuzzyMapping
 
 setting_generator = SettingGenerator()
 setting_updator = SettingUpdator()
@@ -16,6 +17,7 @@ def setting_generator_method(f: callable) -> Callable:
         if node_data.setting_input is None or isinstance(node_data.setting_input, input_schema.NodePromise):
             f(node_data)
         return node_data
+
     setting_generator.add_setting_generator_func(inner)
     return inner
 
@@ -38,37 +40,51 @@ def join(node_data: "NodeData") -> NodeData:
         if len(overlapping_cols) > 0:
             join_key = overlapping_cols[0]
         else:
-            join_key = ''
-        ji = transform_schema.JoinInput(join_mapping=join_key,
-                                        left_select=node_data.main_input.columns,
-                                        right_select=node_data.right_input.columns
-                                        )
-        ji.auto_rename()
-        node_data.setting_input = input_schema.NodeJoin(flow_id=node_data.flow_id,
-                                                        node_id=node_data.node_id,
-                                                        join_input=ji)
+            join_key = ""
+        join_input_manager = transform_schema.JoinInputManager(
+            transform_schema.JoinInput(
+                join_mapping=join_key,
+                left_select=node_data.main_input.columns,
+                right_select=node_data.right_input.columns,
+            )
+        )
+        join_input_manager.auto_rename()
+        ji = join_input_manager.to_join_input()
+        node_data.setting_input = input_schema.NodeJoin(
+            flow_id=node_data.flow_id, node_id=node_data.node_id, join_input=ji
+        )
     return node_data
 
 
 @setting_generator_method
 def cross_join(node_data: "NodeData") -> NodeData:
     if node_data.right_input and node_data.main_input:
-        ji = transform_schema.CrossJoinInput(left_select=node_data.main_input.columns,
-                                             right_select=node_data.right_input.columns)
-        ji.auto_rename()
-        node_data.setting_input = input_schema.NodeCrossJoin(flow_id=node_data.flow_id,
-                                                             node_id=node_data.node_id,
-                                                             cross_join_input=ji)
+        cj_input_manager = transform_schema.CrossJoinInputManager(
+            transform_schema.CrossJoinInput(
+                left_select=node_data.main_input.columns, right_select=node_data.right_input.columns
+            )
+        )
+        cj_input_manager.auto_rename()
+        cj = cj_input_manager.to_cross_join_input()
+        node_data.setting_input = input_schema.NodeCrossJoin(
+            flow_id=node_data.flow_id, node_id=node_data.node_id, cross_join_input=cj
+        )
     return node_data
 
 
 @setting_generator_method
 def filter(node_data: "NodeData") -> NodeData:
     if node_data.main_input:
-        fi = transform_schema.FilterInput(basic_filter=transform_schema.BasicFilter(), filter_type='advanced')
-        node_data.setting_input = input_schema.NodeFilter(flow_id=node_data.flow_id,
-                                                          node_id=node_data.node_id,
-                                                          filter_input=fi)
+        # Default to basic mode with an empty basic filter
+        basic_filter = transform_schema.BasicFilter(
+            field="",
+            operator=transform_schema.FilterOperator.EQUALS,
+            value="",
+        )
+        fi = transform_schema.FilterInput(basic_filter=basic_filter, mode="basic")
+        node_data.setting_input = input_schema.NodeFilter(
+            flow_id=node_data.flow_id, node_id=node_data.node_id, filter_input=fi
+        )
     return node_data
 
 
@@ -80,26 +96,26 @@ def join(node_data: NodeData):
         right_columns = set(node_data.right_input.columns)
         left_select = setting_input.join_input.left_select
         right_select = setting_input.join_input.right_select
+        # Update is_available based on whether column exists in input
         for ls in left_select.renames:
-            if ls.old_name not in right_columns:
-                left_select.remove_select_input(ls.old_name)
+            ls.is_available = ls.old_name in left_columns
         for rs in right_select.renames:
-            if rs.old_name not in right_columns:
-                right_select.remove_select_input(rs.old_name)
-        existing_columns_right = set(r.old_name for r in right_select.renames if r.is_available)
-        existing_columns_left = set(r.old_name for r in left_select.renames if r.is_available)
+            rs.is_available = rs.old_name in right_columns
+        # Check ALL columns in renames to prevent duplicates
+        existing_columns_left = set(r.old_name for r in left_select.renames)
+        existing_columns_right = set(r.old_name for r in right_select.renames)
         missing_incoming_left_columns = [ilc for ilc in left_columns if ilc not in existing_columns_left]
         missing_incoming_right_columns = [irc for irc in right_columns if irc not in existing_columns_right]
-        if not hasattr(setting_input, 'auto_keep_left'):
+        if not hasattr(setting_input, "auto_keep_left"):
             setting_input.auto_keep_left = False
-        if not hasattr(setting_input, 'auto_keep_right'):
+        if not hasattr(setting_input, "auto_keep_right"):
             setting_input.auto_keep_right = False
         for milc in missing_incoming_left_columns:
             select_input = transform_schema.SelectInput(old_name=milc, keep=setting_input.auto_keep_left)
-            setting_input.join_input.add_new_select_column(select_input, 'left')
+            setting_input.join_input.add_new_select_column(select_input, "left")
         for mirc in missing_incoming_right_columns:
             select_input = transform_schema.SelectInput(old_name=mirc, keep=setting_input.auto_keep_right)
-            setting_input.join_input.add_new_select_column(select_input, 'right')
+            setting_input.join_input.add_new_select_column(select_input, "right")
     return node_data
 
 
@@ -111,31 +127,32 @@ def cross_join(node_data: NodeData):
         right_columns = set(node_data.right_input.columns)
         left_select = setting_input.cross_join_input.left_select
         right_select = setting_input.cross_join_input.right_select
+        # Update is_available based on whether column exists in input
         for ls in left_select.renames:
-            if ls.old_name not in right_columns:
-                left_select.remove_select_input(ls.old_name)
+            ls.is_available = ls.old_name in left_columns
         for rs in right_select.renames:
-            if rs.old_name not in right_columns:
-                right_select.remove_select_input(rs.old_name)
-        existing_columns_right = set(r.old_name for r in right_select.renames if r.is_available)
-        existing_columns_left = set(r.old_name for r in left_select.renames if r.is_available)
+            rs.is_available = rs.old_name in right_columns
+        # Check ALL columns in renames to prevent duplicates
+        existing_columns_left = set(r.old_name for r in left_select.renames)
+        existing_columns_right = set(r.old_name for r in right_select.renames)
         missing_incoming_left_columns = [ilc for ilc in left_columns if ilc not in existing_columns_left]
         missing_incoming_right_columns = [irc for irc in right_columns if irc not in existing_columns_right]
-        if not hasattr(setting_input, 'auto_keep_left'):
+        if not hasattr(setting_input, "auto_keep_left"):
             setting_input.auto_keep_left = False
-        if not hasattr(setting_input, 'auto_keep_right'):
+        if not hasattr(setting_input, "auto_keep_right"):
             setting_input.auto_keep_right = False
         for milc in missing_incoming_left_columns:
             select_input = transform_schema.SelectInput(old_name=milc, keep=setting_input.auto_keep_left)
-            setting_input.cross_join_input.add_new_select_column(select_input, 'left')
+            setting_input.cross_join_input.add_new_select_column(select_input, "left")
         for mirc in missing_incoming_right_columns:
             select_input = transform_schema.SelectInput(old_name=mirc, keep=setting_input.auto_keep_right)
-            setting_input.cross_join_input.add_new_select_column(select_input, 'right')
+            setting_input.cross_join_input.add_new_select_column(select_input, "right")
     return node_data
 
 
-def check_if_fuzzy_match_is_valid(left_columns: Iterable[str], right_columns: Iterable[str],
-                                  fuzzy_map: FuzzyMapping) -> bool:
+def check_if_fuzzy_match_is_valid(
+    left_columns: Iterable[str], right_columns: Iterable[str], fuzzy_map: FuzzyMapping
+) -> bool:
     if fuzzy_map.left_col not in left_columns:
         return False
     if fuzzy_map.right_col not in right_columns:
@@ -153,24 +170,24 @@ def fuzzy_match(node_data: NodeData):
         right_select = setting_input.join_input.right_select
         for fuzzy_map in setting_input.join_input.join_mapping:
             fuzzy_map.valid = check_if_fuzzy_match_is_valid(left_columns, right_columns, fuzzy_map)
+        # Update is_available based on whether column exists in input
         for ls in left_select.renames:
-            if ls.old_name not in right_columns:
-                left_select.remove_select_input(ls.old_name)
+            ls.is_available = ls.old_name in left_columns
         for rs in right_select.renames:
-            if rs.old_name not in right_columns:
-                right_select.remove_select_input(rs.old_name)
-        existing_columns_right = set(r.old_name for r in right_select.renames if r.is_available)
-        existing_columns_left = set(r.old_name for r in left_select.renames if r.is_available)
+            rs.is_available = rs.old_name in right_columns
+        # Check ALL columns in renames to prevent duplicates
+        existing_columns_left = set(r.old_name for r in left_select.renames)
+        existing_columns_right = set(r.old_name for r in right_select.renames)
         missing_incoming_left_columns = [ilc for ilc in left_columns if ilc not in existing_columns_left]
         missing_incoming_right_columns = [irc for irc in right_columns if irc not in existing_columns_right]
-        if not hasattr(setting_input, 'auto_keep_left'):
+        if not hasattr(setting_input, "auto_keep_left"):
             setting_input.auto_keep_left = False
-        if not hasattr(setting_input, 'auto_keep_right'):
+        if not hasattr(setting_input, "auto_keep_right"):
             setting_input.auto_keep_right = False
         for milc in missing_incoming_left_columns:
             select_input = transform_schema.SelectInput(old_name=milc, keep=setting_input.auto_keep_left)
-            setting_input.join_input.add_new_select_column(select_input, 'left')
+            setting_input.join_input.add_new_select_column(select_input, "left")
         for mirc in missing_incoming_right_columns:
             select_input = transform_schema.SelectInput(old_name=mirc, keep=setting_input.auto_keep_right)
-            setting_input.join_input.add_new_select_column(select_input, 'right')
+            setting_input.join_input.add_new_select_column(select_input, "right")
     return node_data

flowfile_core/flowfile/sources/external_sources/base_class.py

@@ -1,31 +1,34 @@
 from abc import ABC, abstractmethod
-from typing import Generator, Callable, List, Any, Optional, Dict
-from flowfile_core.flowfile.flow_data_engine.flow_file_column.main import FlowfileColumn
+from collections.abc import Callable, Generator
+from typing import Any
+
 import polars as pl
 
+from flowfile_core.flowfile.flow_data_engine.flow_file_column.main import FlowfileColumn
+
 
 class ExternalDataSource(ABC):
-    schema: Optional[List[FlowfileColumn]]
-    data_getter: Optional[Callable]
+    schema: list[FlowfileColumn] | None
+    data_getter: Callable | None
     is_collected: bool
     cache_store: Any
     _type: str
-    initial_data_getter: Optional[Callable]
+    initial_data_getter: Callable | None
 
     @abstractmethod
     def __init__(self):
         pass
 
     @abstractmethod
-    def get_initial_data(self) -> List[Dict[str, Any]]:
+    def get_initial_data(self) -> list[dict[str, Any]]:
         pass
 
     @abstractmethod
-    def get_iter(self) -> Generator[Dict[str, Any], None, None]:
+    def get_iter(self) -> Generator[dict[str, Any], None, None]:
         pass
 
     @abstractmethod
-    def get_sample(self, n: int = 10000) -> Generator[Dict[str, Any], None, None]:
+    def get_sample(self, n: int = 10000) -> Generator[dict[str, Any], None, None]:
         pass
 
     @abstractmethod
@@ -34,6 +37,5 @@ class ExternalDataSource(ABC):
 
     @staticmethod
     @abstractmethod
-    def parse_schema(*args, **kwargs) -> List[FlowfileColumn]:
+    def parse_schema(*args, **kwargs) -> list[FlowfileColumn]:
         pass
-

flowfile_core/flowfile/sources/external_sources/custom_external_sources/external_source.py

@@ -1,16 +1,19 @@
-from typing import Any, Dict, Generator, List, Optional, Callable
+from collections.abc import Callable, Generator
+from typing import Any
+
+import polars as pl
+
 from flowfile_core.flowfile.flow_data_engine.flow_file_column.main import FlowfileColumn
-from flowfile_core.schemas import input_schema
 from flowfile_core.flowfile.sources.external_sources.base_class import ExternalDataSource
-import polars as pl
+from flowfile_core.schemas import input_schema
 
 
 class CustomExternalSourceSettings:
     data_getter: Generator
-    initial_data_getter: Optional[Callable] = None
-    orientation: str = 'row'
+    initial_data_getter: Callable | None = None
+    orientation: str = "row"
 
-    def __init__(self, data_getter: Generator, initial_data_getter: Optional[Callable] = None, orientation: str = 'row'):
+    def __init__(self, data_getter: Generator, initial_data_getter: Callable | None = None, orientation: str = "row"):
         self.data_getter = data_getter
         self.initial_data_getter = initial_data_getter
         self.orientation = orientation
@@ -18,15 +21,18 @@ class CustomExternalSourceSettings:
 
 class CustomExternalSource(ExternalDataSource):
     data_getter: Generator = None
-    schema: Optional[List[FlowfileColumn]] = None
-    cache_store: List = None
+    schema: list[FlowfileColumn] | None = None
+    cache_store: list = None
     is_collected: bool = False
 
-    def __init__(self, data_getter: Generator[Any, None, None],
-                 initial_data_getter: Callable = None,
-                 orientation: str = 'row',
-                 schema: List = None,
-                 **kwargs):
+    def __init__(
+        self,
+        data_getter: Generator[Any, None, None],
+        initial_data_getter: Callable = None,
+        orientation: str = "row",
+        schema: list = None,
+        **kwargs,
+    ):
         self.cache_store = list()
         self.data_getter = data_getter
         self.collected = False
@@ -38,23 +44,27 @@ class CustomExternalSource(ExternalDataSource):
         else:
             self.schema = None
 
-        if not initial_data_getter and orientation == 'row':
+        if not initial_data_getter and orientation == "row":
+
             def initial_data_getter():
                 if len(self.cache_store) == 0:
                     self.cache_store.append(next(data_getter, None))
                 return self.cache_store
+
             self.initial_data_getter = initial_data_getter
         elif initial_data_getter:
             self.initial_data_getter = initial_data_getter
         elif self.schema:
+
             def initial_data_getter():
                 return [{d.column_name: None for d in self.schema}]
+
             self.initial_data_getter = initial_data_getter
         else:
             self.initial_data_getter = None
 
     @staticmethod
-    def parse_schema(schema: List[Any]) -> List[FlowfileColumn]:
+    def parse_schema(schema: list[Any]) -> list[FlowfileColumn]:
         if len(schema) == 0:
             return []
         first_col = schema[0]
@@ -63,7 +73,7 @@ class CustomExternalSource(ExternalDataSource):
         elif isinstance(first_col, (list, tuple)):
             return [FlowfileColumn.from_input(column_name=col[0], data_type=col[1]) for col in schema]
         elif isinstance(first_col, str):
-            return [FlowfileColumn.from_input(column_name=col, data_type='varchar') for col in schema]
+            return [FlowfileColumn.from_input(column_name=col, data_type="varchar") for col in schema]
         elif isinstance(first_col, input_schema.MinimalFieldInfo):
             return [FlowfileColumn.from_input(column_name=col.name, data_type=col.data_type) for col in schema]
         elif isinstance(first_col, FlowfileColumn):
@@ -76,7 +86,7 @@ class CustomExternalSource(ExternalDataSource):
             return self.initial_data_getter()
         return []
 
-    def get_iter(self) -> Generator[Dict[str, Any], None, None]:
+    def get_iter(self) -> Generator[dict[str, Any], None, None]:
         if self.collected:
             return
         for data in self.cache_store:

flowfile_core/flowfile/sources/external_sources/custom_external_sources/sample_users.py

@@ -1,10 +1,13 @@
-from typing import Dict, Any, Generator
+from collections.abc import Generator
 from time import sleep
-from flowfile_core.schemas.input_schema import SampleUsers
+from typing import Any
+
 import requests
 
+from flowfile_core.schemas.input_schema import SampleUsers
 
-def getter(data: SampleUsers) -> Generator[Dict[str, Any], None, None]:
+
+def getter(data: SampleUsers) -> Generator[dict[str, Any], None, None]:
     """
     Sample users generator function. This is a minimal example of a generator function that yields user data and can
     be used in a flowfile. The function simulates a delay to mimic the behavior of an external data source.
@@ -17,13 +20,10 @@ def getter(data: SampleUsers) -> Generator[Dict[str, Any], None, None]:
     index_pos = 0
     for i in range(data.size):
         sleep(0.01)
-        headers = {
-            'x-api-key': 'reqres-free-v1'
-        }
+        headers = {"x-api-key": "reqres-free-v1"}
 
         response = requests.get("https://reqres.in/api/users", headers=headers).json()
-        for v in response['data']:
-            v['index'] = index_pos
+        for v in response["data"]:
+            v["index"] = index_pos
             index_pos += 1
             yield v
-

flowfile_core/flowfile/sources/external_sources/factory.py

@@ -16,4 +16,3 @@ def data_source_factory(source_type: str, **kwargs) -> CustomExternalSource:
         return CustomExternalSource(**kwargs)
     else:
         raise ValueError(f"Unknown source type: {source_type}")
-