yaml / pyyaml

Canonical source repository for PyYAML
MIT License
2.47k stars 507 forks source link

Support for `enum.StrEnum` #722

Open trim21 opened 1 year ago

trim21 commented 1 year ago

hope pyyaml can support enum.StrEnum

element of StrEnum is consider str in Python

import enum

import yaml

class Status(enum.StrEnum):
    A = 'a'
    B = 'b'

print(isinstance(Status.A, str)) # True
print(yaml.safe_dump({'status': Status.A}))
Traceback (most recent call last):
  File "C:\Users\Trim21\proj\test\a.py", line 12, in <module>
    print(yaml.safe_dump({'status': Status.A}))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Trim21\proj\test\.venv\Lib\site-packages\yaml\__init__.py", line 269, in safe_dump
    return dump_all([data], stream, Dumper=SafeDumper, **kwds)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Trim21\proj\test\.venv\Lib\site-packages\yaml\__init__.py", line 241, in dump_all
    dumper.represent(data)
  File "C:\Users\Trim21\proj\test\.venv\Lib\site-packages\yaml\representer.py", line 27, in represent
    node = self.represent_data(data)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Trim21\proj\test\.venv\Lib\site-packages\yaml\representer.py", line 48, in represent_data
    node = self.yaml_representers[data_types[0]](self, data)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Trim21\proj\test\.venv\Lib\site-packages\yaml\representer.py", line 207, in represent_dict
    return self.represent_mapping('tag:yaml.org,2002:map', data)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Trim21\proj\test\.venv\Lib\site-packages\yaml\representer.py", line 118, in represent_mapping
    node_value = self.represent_data(item_value)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Trim21\proj\test\.venv\Lib\site-packages\yaml\representer.py", line 58, in represent_data
    node = self.yaml_representers[None](self, data)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Trim21\proj\test\.venv\Lib\site-packages\yaml\representer.py", line 231, in represent_undefined
    raise RepresenterError("cannot represent an object", data)
yaml.representer.RepresenterError: ('cannot represent an object', <Status.A: 'a'>)

stdlib json support this: print(json.dumps({'status': Status.A}))

spanezz commented 4 months ago

Current workaround for this specific case:

  yaml.SafeDumper.add_representer(
      Status,
      yaml.representer.SafeRepresenter.represent_str,
  )                                                            

It may need to be repeated for the all the various dumpers that may be used used, and it unfortunately does not work with StrEnum instead of Status

bergtholdt commented 4 months ago

Just use add_multi_representer for the base class, so all sub-classes will use the base representer:

  yaml.SafeDumper.add_multi_representer(
      StrEnum,
      yaml.representer.SafeRepresenter.represent_str,
  )
trim21 commented 4 months ago

Just use add_multi_representer for the base class, so all sub-classes will use the base representer:

  yaml.SafeDumper.add_multi_representer(
      StrEnum,
      yaml.representer.SafeRepresenter.represent_str,
  )

would be nice to have this included in pyyaml library