Repository: materialsproject/fireworks

Issue:
AttributeError: "safe_load()" has been removed with ruamel.yaml 0.18.5
Please accept my apologies for bypassing https://discuss.matsci.org/c/fireworks but what is below is obviously an issue and no help request.

I have a fresh fireworks install in a fresh virtual environment (either conda or venv) with python 3.10 (python==3.10.12, ruamel.yaml==0.18.5, FireWorks==2.0.3) or newer python, I have issues with any use of fireworks related to a removed function in the ruamel.yaml package. For example, a simple `lpad get_wflows -d count` yields this error message:

```
Traceback (most recent call last):
  File "test/lib/python3.10/site-packages/fireworks/scripts/lpad_run.py", line 108, in get_lp
    lp = LaunchPad.from_file(args.launchpad_file)
  File "test/lib/python3.10/site-packages/fireworks/utilities/fw_serializers.py", line 292, in from_file
    return cls.from_format(f.read(), f_format=f_format)
  File "test/lib/python3.10/site-packages/fireworks/utilities/fw_serializers.py", line 257, in from_format
    dct = yaml.safe_load(f_str)
  File "test/lib/python3.10/site-packages/ruamel/yaml/main.py", line 1105, in safe_load
    error_deprecation('safe_load', 'load', arg="typ='safe', pure=True")
  File "test/lib/python3.10/site-packages/ruamel/yaml/main.py", line 1039, in error_deprecation
    raise AttributeError(s, name=None)
AttributeError:
"safe_load()" has been removed, use

  yaml = YAML(typ='safe', pure=True)
  yaml.load(...)

instead of file "test/lib/python3.10/site-packages/fireworks/utilities/fw_serializers.py", line 257

            dct = yaml.safe_load(f_str)



During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "test/bin/lpad", line 8, in <module>
    sys.exit(lpad())
  File "test/lib/python3.10/site-packages/fireworks/scripts/lpad_run.py", line 1555, in lpad
    args.func(args)
  File "test/lib/python3.10/site-packages/fireworks/scripts/lpad_run.py", line 421, in get_wfs
    lp = get_lp(args)
  File "test/lib/python3.10/site-packages/fireworks/scripts/lpad_run.py", line 122, in get_lp
    f"FireWorks was not able to connect to MongoDB at {lp.host}:{lp.port}. Is the server running? "
UnboundLocalError: local variable 'lp' referenced before assignment
```

Using both the latest release (2.0.3) and the main branch results in the same behavior.

The workaround is to downgrade ruamel.yaml manually that is unfortunately an insufficient solution for the users of our software that uses fireworks as backend.

Hints from issue comments:
@ikondov: Just to complete the picture, with ruamel 0.17.40 I get this warning:

```
  test-venv/lib/python3.10/site-packages/fireworks/utilities/fw_serializers.py:256: PendingDeprecationWarning:
  safe_load will be removed, use

    yaml=YAML(typ='safe', pure=True)
    yaml.load(...)

  instead
    dct = yaml.safe_load(f_str)
```

Solution patch:
diff --git a/fireworks/scripts/lpad_run.py b/fireworks/scripts/lpad_run.py
index b6748dda..46ebb65b 100644
--- a/fireworks/scripts/lpad_run.py
+++ b/fireworks/scripts/lpad_run.py
@@ -6,12 +6,14 @@ import datetime
 import json
 import os
 import re
+import sys
 import time
 from argparse import ArgumentParser, ArgumentTypeError, Namespace
 from importlib import metadata
 from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Union
 
-import ruamel.yaml as yaml
+from ruamel.yaml import YAML
+
 from pymongo import ASCENDING, DESCENDING
 
 from fireworks import FW_INSTALL_DIR
@@ -260,8 +262,7 @@ def print_fws(ids, lp, args: Namespace) -> None:
             fws.append(d)
     if len(fws) == 1:
         fws = fws[0]
-
-    print(args.output(fws))
+    get_output(args, fws)
 
 
 def get_fw_ids_helper(lp: LaunchPad, args: Namespace, count_only: Union[bool, None] = None) -> Union[List[int], int]:
@@ -343,7 +344,7 @@ def get_fws(args: Namespace) -> None:
     lp = get_lp(args)
     ids = get_fw_ids_helper(lp, args)
     fws = get_fws_helper(lp, ids, args)
-    print(args.output(fws))
+    get_output(args, fws)
 
 
 def get_fws_in_wfs(args: Namespace) -> None:
@@ -467,7 +468,7 @@ def get_wfs(args: Namespace) -> None:
     else:
         if len(wfs) == 1:
             wfs = wfs[0]
-        print(args.output(wfs))
+        get_output(args, wfs)
 
 
 def delete_wfs(args: Namespace) -> None:
@@ -852,13 +853,18 @@ def orphaned(args: Namespace) -> None:
         lp.m_logger.info(f"Found {len(orphaned_fw_ids)} orphaned fw_ids: {orphaned_fw_ids}")
         lp.delete_fws(orphaned_fw_ids, delete_launch_dirs=args.delete_launch_dirs)
     else:
-        print(args.output(fws))
+        get_output(args, fws)
 
 
-def get_output_func(format: Literal["json", "yaml"]) -> Callable[[str], Any]:
-    if format == "json":
-        return lambda x: json.dumps(x, default=DATETIME_HANDLER, indent=4)
-    return lambda x: yaml.safe_dump(recursive_dict(x, preserve_unicode=False), default_flow_style=False)
+def get_output(args: Namespace, objs: List[Any]) -> None:
+    """prints output on stdout"""
+    if args.output == "json":
+        json.dump(objs, sys.stdout, default=DATETIME_HANDLER, indent=4)
+    else:
+        yaml = YAML(typ='safe', pure=True)
+        yaml.default_flow_style = False
+        yaml.dump(recursive_dict(objs, preserve_unicode=False), sys.stdout)
+    print()
 
 
 def arg_positive_int(value: str) -> int:
@@ -1539,8 +1545,6 @@ def lpad(argv: Optional[Sequence[str]] = None) -> int:
         cfg_files_to_check.append(("fworker", "-w", False, FWORKER_LOC))
     _validate_config_file_paths(args, cfg_files_to_check)
 
-    args.output = get_output_func(args.output)
-
     if args.command is None:
         # if no command supplied, print help
         parser.print_help()
diff --git a/fireworks/user_objects/firetasks/dataflow_tasks.py b/fireworks/user_objects/firetasks/dataflow_tasks.py
index e2a8d499..7103b30a 100644
--- a/fireworks/user_objects/firetasks/dataflow_tasks.py
+++ b/fireworks/user_objects/firetasks/dataflow_tasks.py
@@ -6,6 +6,8 @@ __copyright__ = "Copyright 2016, Karlsruhe Institute of Technology"
 
 import sys
 
+from ruamel.yaml import YAML
+
 from fireworks import Firework
 from fireworks.core.firework import FiretaskBase, FWAction
 from fireworks.utilities.fw_serializers import load_object
@@ -300,7 +302,7 @@ class ForeachTask(FiretaskBase):
         chunklen = lensplit // nchunks
         if lensplit % nchunks > 0:
             chunklen = chunklen + 1
-        chunks = [split_field[i : i + chunklen] for i in range(0, lensplit, chunklen)]
+        chunks = [split_field[i:i+chunklen] for i in range(0, lensplit, chunklen)]
 
         fireworks = []
         for index, chunk in enumerate(chunks):
@@ -381,8 +383,6 @@ class ImportDataTask(FiretaskBase):
         import operator
         from functools import reduce
 
-        import ruamel.yaml as yaml
-
         filename = self["filename"]
         mapstring = self["mapstring"]
         assert isinstance(filename, basestring)
@@ -392,7 +392,7 @@ class ImportDataTask(FiretaskBase):
         fmt = filename.split(".")[-1]
         assert fmt in ["json", "yaml"]
         with open(filename) as inp:
-            data = json.load(inp) if fmt == "json" else yaml.safe_load(inp)
+            data = json.load(inp) if fmt == "json" else YAML(typ="safe", pure=True).load(inp)
 
         leaf = reduce(operator.getitem, maplist[:-1], fw_spec)
         if isinstance(data, dict):
diff --git a/fireworks/utilities/fw_serializers.py b/fireworks/utilities/fw_serializers.py
index 36c879ba..e9e89b32 100644
--- a/fireworks/utilities/fw_serializers.py
+++ b/fireworks/utilities/fw_serializers.py
@@ -25,7 +25,6 @@ Some advantages:
         their fw_name (if no parameters are needed to describe the object)
 
 """
-
 import abc
 import datetime
 import importlib
@@ -34,7 +33,8 @@ import json  # note that ujson is faster, but at this time does not support "def
 import pkgutil
 import traceback
 
-import ruamel.yaml as yaml
+from ruamel.yaml import YAML
+from ruamel.yaml.compat import StringIO
 from monty.json import MontyDecoder, MSONable
 
 from fireworks.fw_config import (
@@ -235,8 +235,12 @@ class FWSerializable(metaclass=abc.ABCMeta):
         if f_format == "json":
             return json.dumps(self.to_dict(), default=DATETIME_HANDLER, **kwargs)
         if f_format == "yaml":
-            # start with the JSON format, and convert to YAML
-            return yaml.safe_dump(self.to_dict(), default_flow_style=YAML_STYLE, allow_unicode=True)
+            yaml = YAML(typ="safe", pure=True)
+            yaml.default_flow_style = YAML_STYLE
+            yaml.allow_unicode = True
+            strm = StringIO()
+            yaml.dump(self.to_dict(), strm)
+            return strm.getvalue()
         raise ValueError(f"Unsupported format {f_format}")
 
     @classmethod
@@ -254,7 +258,7 @@ class FWSerializable(metaclass=abc.ABCMeta):
         if f_format == "json":
             dct = json.loads(f_str)
         elif f_format == "yaml":
-            dct = yaml.safe_load(f_str)
+            dct = YAML(typ="safe", pure=True).load(f_str)
         else:
             raise ValueError(f"Unsupported format {f_format}")
         if JSON_SCHEMA_VALIDATE and cls.__name__ in JSON_SCHEMA_VALIDATE_LIST:
@@ -271,8 +275,16 @@ class FWSerializable(metaclass=abc.ABCMeta):
         """
         if f_format is None:
             f_format = filename.split(".")[-1]
-        with open(filename, "w", **ENCODING_PARAMS) as f:
-            f.write(self.to_format(f_format=f_format, **kwargs))
+        with open(filename, "w", **ENCODING_PARAMS) as f_out:
+            if f_format == "json":
+                json.dump(self.to_dict(), f_out, default=DATETIME_HANDLER, **kwargs)
+            elif f_format == "yaml":
+                yaml = YAML(typ="safe", pure=True)
+                yaml.default_flow_style = YAML_STYLE
+                yaml.allow_unicode = True
+                yaml.dump(self.to_dict(), f_out)
+            else:
+                raise ValueError(f"Unsupported format {f_format}")
 
     @classmethod
     def from_file(cls, filename, f_format=None):
@@ -387,7 +399,7 @@ def load_object_from_file(filename, f_format=None):
         if f_format == "json":
             dct = json.loads(f.read())
         elif f_format == "yaml":
-            dct = yaml.safe_load(f)
+            dct = YAML(typ="safe", pure=True).load(f.read())
         else:
             raise ValueError(f"Unknown file format {f_format} cannot be loaded!")
 


Test patch:
diff --git a/fireworks/user_objects/firetasks/tests/test_dataflow_tasks.py b/fireworks/user_objects/firetasks/tests/test_dataflow_tasks.py
index 58932d5b..45404376 100644
--- a/fireworks/user_objects/firetasks/tests/test_dataflow_tasks.py
+++ b/fireworks/user_objects/firetasks/tests/test_dataflow_tasks.py
@@ -5,6 +5,8 @@ import unittest
 import uuid
 from unittest import SkipTest
 
+from ruamel.yaml import YAML
+
 from fireworks.user_objects.firetasks.dataflow_tasks import (
     CommandLineTask,
     ForeachTask,
@@ -293,11 +295,9 @@ class ImportDataTaskTest(unittest.TestCase):
         """Loads data from a file into spec."""
         import json
 
-        import ruamel.yaml as yaml
-
         temperature = {"value": 273.15, "units": "Kelvin"}
         spec = {"state parameters": {}}
-        formats = {"json": json, "yaml": yaml}
+        formats = {"json": json, "yaml": YAML(typ="safe", pure=True)}
         params = {"mapstring": "state parameters/temperature"}
         for fmt in formats:
             filename = str(uuid.uuid4()) + "." + fmt
diff --git a/fireworks/user_objects/queue_adapters/tests/test_common_adapter.py b/fireworks/user_objects/queue_adapters/tests/test_common_adapter.py
index c08d0289..bbb33ef1 100644
--- a/fireworks/user_objects/queue_adapters/tests/test_common_adapter.py
+++ b/fireworks/user_objects/queue_adapters/tests/test_common_adapter.py
@@ -7,8 +7,11 @@ __maintainer__ = "Shyue Ping Ong"
 __email__ = "shyuep@gmail.com"
 __date__ = "12/31/13"
 
+import sys
 import unittest
 
+from ruamel.yaml import YAML
+
 from fireworks.user_objects.queue_adapters.common_adapter import CommonAdapter, os
 from fireworks.utilities.fw_serializers import load_object, load_object_from_file
 
@@ -41,9 +44,10 @@ class CommonAdapterTest(unittest.TestCase):
         p = load_object_from_file(os.path.join(os.path.dirname(__file__), "pbs.yaml"))
         p = CommonAdapter(q_type="PBS", q_name="hello", ppnode="8:ib", nnodes=1, hello="world", queue="random")
         print(p.get_script_str("."))
-        import ruamel.yaml as yaml
-
-        print(yaml.safe_dump(p.to_dict(), default_flow_style=False))
+        yaml = YAML(typ="safe", pure=True)
+        yaml.default_flow_style = False
+        yaml.dump(p.to_dict(), sys.stdout)
+        print()
 
     def test_parse_njobs(self):
         pbs = """