Merge pull request #628 from wrangleworks/yaml-CLoader

Refactor YAML handling to use CSafeDumper and CSafeLoader for improved performance
wrangleworks · Feb 24, 2025 · e836b9b · e836b9b
2 parents 52c941b + 0afa7ad
commit e836b9b
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 5 deletions.
diff --git a/wrangles/openai.py b/wrangles/openai.py
@@ -7,6 +7,11 @@
 import requests as _requests
 import numpy as _np
 import time as _time
+try:
+    from yaml import CSafeDumper as _YAMLDumper
+except ImportError:
+    from yaml import SafeDumper as _YAMLDumper
+
 
 def chatGPT(
     data: any,
@@ -28,7 +33,14 @@ def chatGPT(
     if len(data) == 1:
         content = list(data.values())[0]
     else:
-        content = _yaml.dump(data, indent=2, sort_keys=False, allow_unicode=True)
+        content = _yaml.dump(
+            data,
+            indent=2,
+            sort_keys=False,
+            allow_unicode=True,
+            Dumper=_YAMLDumper,
+            width=1000
+        )
 
     settings_local = _copy.deepcopy(settings)
     settings_local["messages"].append(

diff --git a/wrangles/recipe.py b/wrangles/recipe.py
@@ -30,7 +30,10 @@
     add_special_parameters as _add_special_parameters,
     wildcard_expansion as _wildcard_expansion
 )
-
+try:
+    from yaml import CSafeLoader as _YamlLoader, CSafeDumper as _YAMLDumper
+except ImportError:
+    from yaml import SafeLoader as _YamlLoader, SafeDumper as _YAMLDumper
 
 _logging.getLogger().setLevel(_logging.INFO)
 
@@ -122,7 +125,7 @@ def _replace_templated_values(
                 and '\n' in replacement_value
             ):
                 try:
-                    replacement_value = _yaml.safe_load(replacement_value)
+                    replacement_value = _yaml.load(replacement_value, Loader=_YamlLoader)
                 except:
                     # Replacement wasn't YAML
                     pass
@@ -175,7 +178,7 @@ def _load_recipe(
     if not isinstance(recipe, str):
         try:
             # If user passes in a pre-parsed recipe, convert back to YAML
-            recipe = _yaml.dump(recipe, sort_keys=False, allow_unicode=True)
+            recipe = _yaml.dump(recipe, sort_keys=False, Dumper=_YAMLDumper, allow_unicode=True)
         except:
             raise ValueError('Recipe passed in as an invalid type')
 

diff --git a/wrangles/recipe_wrangles/convert.py b/wrangles/recipe_wrangles/convert.py
@@ -8,6 +8,10 @@
 import pandas as _pd
 from fractions import Fraction as _Fraction
 import yaml as _yaml
+try:
+    from yaml import CSafeLoader as _YAMLLoader, CSafeDumper as _YAMLDumper
+except ImportError:
+    from yaml import SafeLoader as _YAMLLoader, SafeDumper as _YAMLDumper
 
 
 def case(df: _pd.DataFrame, input: _Union[str, list], output: _Union[str, list] = None, case: str = 'lower') -> _pd.DataFrame:
@@ -374,7 +378,7 @@ def _load_with_fallback(value):
         If no default, raise an error.
         """
         try:
-            return _yaml.safe_load(value, **kwargs) or default
+            return _yaml.load(value, Loader=_YAMLLoader, **kwargs) or default
         except:
             if default != None:
                 return default
@@ -458,6 +462,7 @@ def to_yaml(
                     row,
                     sort_keys=sort_keys,
                     allow_unicode=allow_unicode,
+                    Dumper=_YAMLDumper,
                     **kwargs
                 )
                 for row in df[input_columns].values
@@ -470,6 +475,7 @@ def to_yaml(
                 row,
                 sort_keys=sort_keys,
                 allow_unicode=allow_unicode,
+                Dumper=_YAMLDumper,
                 **kwargs
             )
             for row in df[input].to_dict(orient="records")