Fixes serialize for extract fields

Author: elijahbenizzy

examples/parallelism/file_processing/run.py

@@ -1,3 +1,9 @@
+import warnings
+
+from hamilton.lifecycle import SlowDownYouMoveTooFast
+
+warnings.simplefilter(action="ignore", category=FutureWarning)
+
 import logging
 
 import aggregate_data
@@ -9,9 +15,9 @@
 
 from hamilton import driver, log_setup
 from hamilton.execution import executors
-from hamilton.plugins import h_dask, h_ray
+from hamilton.plugins import h_dask, h_ray, h_rich
 
-log_setup.setup_logging(logging.INFO)
+log_setup.setup_logging(logging.FATAL)
 
 
 @click.command()
@@ -39,6 +45,10 @@ def main(mode: str):
         driver.Builder()
         .enable_dynamic_execution(allow_experimental_mode=True)
         .with_remote_executor(remote_executor)  # We only need to specify remote exeecutor
+        .with_adapters(
+            h_rich.RichProgressBar(),
+            SlowDownYouMoveTooFast(sleep_time_mean=0.3, sleep_time_std=0.1),
+        )
         # The local executor just runs it synchronously
         .with_modules(aggregate_data, list_data, process_data)
         .build()

hamilton/function_modifiers/expanders.py

@@ -731,6 +731,36 @@ def _validate_extract_fields(fields: dict):
             )
 
 
+async def dict_generator_async(
+    *args,
+    fn,
+    fill_with,
+    fields,
+    **kwargs,
+):
+    dict_generated = await fn(*args, **kwargs)
+    if fill_with is not None:
+        for field in fields:
+            if field not in dict_generated:
+                dict_generated[field] = fill_with
+    return dict_generated
+
+
+async def dict_generator(
+    *args,
+    fn,
+    fill_with,
+    fields,
+    **kwargs,
+):
+    dict_generated = fn(*args, **kwargs)
+    if fill_with is not None:
+        for field in fields:
+            if field not in dict_generated:
+                dict_generated[field] = fill_with
+    return dict_generated
+
+
 class extract_fields(base.SingleNodeNodeTransformer):
     """Extracts fields from a dictionary of output."""
 
@@ -804,29 +834,35 @@ def transform_node(
         """
         fn = node_.callable
         base_doc = node_.documentation
-
+        dict_generator_fn = (
+            functools.partial(dict_generator, fn=fn, fill_with=self.fill_with, fields=self.fields)
+            if not (inspect.iscoroutinefunction(fn))
+            else functools.partial(
+                dict_generator_async, fn=fn, fill_with=self.fill_with, fields=self.fields
+            )
+        )
         # if fn is async
-        if inspect.iscoroutinefunction(fn):
-
-            async def dict_generator(*args, **kwargs):
-                dict_generated = await fn(*args, **kwargs)
-                if self.fill_with is not None:
-                    for field in self.fields:
-                        if field not in dict_generated:
-                            dict_generated[field] = self.fill_with
-                return dict_generated
-
-        else:
-
-            def dict_generator(*args, **kwargs):
-                dict_generated = fn(*args, **kwargs)
-                if self.fill_with is not None:
-                    for field in self.fields:
-                        if field not in dict_generated:
-                            dict_generated[field] = self.fill_with
-                return dict_generated
-
-        output_nodes = [node_.copy_with(callabl=dict_generator)]
+        # if inspect.iscoroutinefunction(fn):
+        #
+        #     async def dict_generator(*args, **kwargs):
+        #         dict_generated = await fn(*args, **kwargs)
+        #         if self.fill_with is not None:
+        #             for field in self.fields:
+        #                 if field not in dict_generated:
+        #                     dict_generated[field] = self.fill_with
+        #         return dict_generated
+        #
+        # else:
+        #
+        #     def dict_generator(*args, **kwargs):
+        #         dict_generated = fn(*args, **kwargs)
+        #         if self.fill_with is not None:
+        #             for field in self.fields:
+        #                 if field not in dict_generated:
+        #                     dict_generated[field] = self.fill_with
+        #         return dict_generated
+
+        output_nodes = [node_.copy_with(callabl=dict_generator_fn)]
 
         for field, field_type in self.fields.items():
             doc_string = base_doc  # default doc string of base function.

pyproject.toml

@@ -57,8 +57,7 @@ docs = [
   "diskcache",
   # required for all the plugins
   "dlt",
-  # furo -- install from main for now until the next release is out:
-  "furo @ git+https://github.com/pradyunsg/furo@main",
+  "furo",
   "gitpython", # Required for parsing git info for generation of data-adapter docs
   "grpcio-status",
   "lightgbm",