[YamlWriter] Set UnitSystem directly in YamlWriter::setUnits

include/cantera/base/YamlWriter.h

@@ -59,6 +59,12 @@ class YamlWriter
     //!     corresponding units supported by the UnitSystem class.
     void setUnits(const std::map<std::string, std::string>& units={});
 
+    //! Set the units to be used in the output file. Dimensions not specified
+    //! will use Cantera's defaults.
+    //! @param units  A UnitSystem object specifying dimensions (mass, length, time,
+    //!     quantity, pressure, energy, activation-energy).
+    void setUnits(const UnitSystem& units=UnitSystem());
+
 protected:
     std::vector<shared_ptr<Solution>> m_phases;
 

interfaces/cython/cantera/_cantera.pxd

@@ -649,6 +649,7 @@ cdef extern from "cantera/base/YamlWriter.h" namespace "Cantera":
         void setPrecision(int)
         void skipUserDefined(cbool)
         void setUnits(stdmap[string, string]&) except +translate_exception
+        void setUnits(CxxUnitSystem) except +translate_exception
 
 cdef extern from "cantera/equil/MultiPhase.h" namespace "Cantera":
     cdef cppclass CxxMultiPhase "Cantera::MultiPhase":
@@ -1223,6 +1224,8 @@ cdef class DustyGasTransport(Transport):
 cdef class YamlWriter:
     cdef shared_ptr[CxxYamlWriter] _writer
     cdef CxxYamlWriter* writer
+    @staticmethod
+    cdef CxxUnitSystem _get_unitsystem(UnitSystem units)
 
 cdef class Mixture:
     cdef CxxMultiPhase* mix

interfaces/cython/cantera/base.pyx

@@ -263,8 +263,8 @@ cdef class _SolutionBase:
             Additional ThermoPhase / Solution objects to be included in the
             output file
         :param units:
-            A dictionary of the units to be used for each dimension. See
-            `YamlWriter.output_units`.
+            A `UnitSystem` object or dictionary of the units to be used for
+            each dimension. See `YamlWriter.output_units`.
         :param precision:
             For output floating point values, the maximum number of digits to
             the right of the decimal point. The default is 15 digits.

interfaces/cython/cantera/yamlwriter.pyx

@@ -50,19 +50,20 @@ cdef class YamlWriter:
         will use Cantera's defaults.
 
         :param units:
-            A UnitSystem object or map where keys are dimensions (mass, length, time,
+            A `UnitSystem` object or map where keys are dimensions (mass, length, time,
             quantity, pressure, energy, activation-energy), and the values are
             corresponding units such as kg, mm, s, kmol, Pa, cal, and eV.
         """
         def __set__(self, units):
-            if isinstance(units, UnitSystem):
-                defaults = UnitSystem().units
-                units = {k: v for k, v in units.units.items() if defaults[k] != v}
-            cdef stdmap[string, string] cxxunits
-            for dimension, unit in units.items():
-                cxxunits[stringify(dimension)] = stringify(unit)
+            if not isinstance(units, UnitSystem):
+                units = UnitSystem(units)
+            cdef CxxUnitSystem cxxunits = YamlWriter._get_unitsystem(units)
             self.writer.setUnits(cxxunits)
 
+    @staticmethod
+    cdef CxxUnitSystem _get_unitsystem(UnitSystem units):
+        return units.unitsystem
+
     def __reduce__(self):
         raise NotImplementedError('YamlWriter object is not picklable')
 

src/base/YamlWriter.cpp

@@ -160,4 +160,9 @@ void YamlWriter::setUnits(const std::map<std::string, std::string>& units)
     m_output_units.setDefaults(units);
 }
 
+void YamlWriter::setUnits(const UnitSystem& units)
+{
+    m_output_units = units;
+}
+
 }