pint icon indicating copy to clipboard operation
pint copied to clipboard

Unpickling Custom Quantities/Units doesn't use custom unit registry

Open zobac opened this issue 2 years ago • 2 comments

I've updated my custom Quantity, Unit, and UREG classes as per the docs after you helped me with my last issue. Once done, things work great until I need to pickle and unpickle my data.

It looks like the pint unpickler doesn't consider the custom UnitRegistry, which results in pint.errors.UndefinedUnitError:


import pint
class MyQuantity(pint.UnitRegistry.Quantity):

    # Notice that subclassing pint.Quantity
    # is not necessary.
    # Pint will inspect the Registry class and create
    # a Quantity class that contains all the
    # required parents.

    def to_my_desired_format(self):
        """Do something else
        """
class MyUnit(pint.UnitRegistry.Unit):

    # Notice that subclassing pint.Quantity
    # is not necessary.
    # Pint will inspect the Registry class and create
    # a Quantity class that contains all the
    # required parents.

    def to_my_desired_format(self):
        """Do something else
        """

from typing_extensions import TypeAlias 
class MyRegistry(pint.registry.GenericUnitRegistry[MyQuantity, pint.Unit]):

    Quantity: TypeAlias = MyQuantity
    Unit: TypeAlias = MyUnit

ureg = MyRegistry()
ureg.define('dog_year = 52 * day = dy')
lassie_lifespan = ureg.Quantity(10, 'dog_years')
p = pickle.dumps(lassie_lifespan)
up = pickle.loads(p)
Traceback (most recent call last):
  Python Shell, prompt 9, line 1
    # Used internally for debug sandbox under external interpreter
  File "d:\envs\xarray\lib\site-packages\pint\__init__.py", line 85, in _unpickle_quantity
    return _unpickle(application_registry.Quantity, *args)
  File "d:\envs\xarray\lib\site-packages\pint\__init__.py", line 78, in _unpickle
    application_registry.parse_units(name)
  File "d:\envs\xarray\lib\site-packages\pint\facets\plain\registry.py", line 1127, in parse_units
    units = self._parse_units(input_string, as_delta, case_sensitive)
  File "d:\envs\xarray\lib\site-packages\pint\facets\nonmultiplicative\registry.py", line 70, in _parse_units
    return super()._parse_units(input_string, as_delta, case_sensitive)
  File "d:\envs\xarray\lib\site-packages\pint\facets\plain\registry.py", line 1160, in _parse_units
    cname = self.get_name(name, case_sensitive=case_sensitive)
  File "d:\envs\xarray\lib\site-packages\pint\facets\plain\registry.py", line 619, in get_name
    raise UndefinedUnitError(name_or_alias)
pint.errors.UndefinedUnitError: 'dog_year' is not defined in the unit registry

zobac avatar Jul 07 '23 20:07 zobac

I should note that the same behavior occurs if I make pint use my unit registry:

ureg = MyRegistry()
pint.UnitRegistry = lambda : ureg
ureg.define('dog_year = 52 * day = dy')
lassie_lifespan = ureg.Quantity(10, 'dog_years')
p = pickle.dumps(lassie_lifespan)
up = pickle.loads(p)

Traceback (most recent call last):
  Python Shell, prompt 7, line 2
    if __name__ == '__main__':
  File "d:\envs\xarray\lib\site-packages\pint\__init__.py", line 85, in _unpickle_quantity
    return _unpickle(application_registry.Quantity, *args)
  File "d:\envs\xarray\lib\site-packages\pint\__init__.py", line 78, in _unpickle
    application_registry.parse_units(name)
  File "d:\envs\xarray\lib\site-packages\pint\facets\plain\registry.py", line 1127, in parse_units
    units = self._parse_units(input_string, as_delta, case_sensitive)
  File "d:\envs\xarray\lib\site-packages\pint\facets\nonmultiplicative\registry.py", line 70, in _parse_units
    return super()._parse_units(input_string, as_delta, case_sensitive)
  File "d:\envs\xarray\lib\site-packages\pint\facets\plain\registry.py", line 1160, in _parse_units
    cname = self.get_name(name, case_sensitive=case_sensitive)
  File "d:\envs\xarray\lib\site-packages\pint\facets\plain\registry.py", line 619, in get_name
    raise UndefinedUnitError(name_or_alias)
pint.errors.UndefinedUnitError: 'dog_year' is not defined in the unit registry
``

zobac avatar Jul 17 '23 13:07 zobac

I guess for now I'll create my own _unpickle_quantity and _unpickle_unit functions and punch them in to pint.

def _unpickle(cls, *args):
    """Rebuild object upon unpickling.
    All units must exist in the application registry.

    Parameters
    ----------
    cls : Quantity, Magnitude, or Unit
    *args

    Returns
    -------
    object of type cls

    """
    from pint.util import UnitsContainer

    for arg in args:
        # Prefixed units are defined within the registry
        # on parsing (which does not happen here).
        # We need to make sure that this happens before using.
        if isinstance(arg, UnitsContainer):
            for name in arg:
                ureg.parse_units(name)

    return cls(*args)



def _unpickle_quantity(cls, *args):
    """Rebuild quantity upon unpickling using the application registry."""
    return _unpickle(ureg.Quantity, *args)


def _unpickle_unit(cls, *args):
    """Rebuild unit upon unpickling using the application registry."""
    return _unpickle(ureg.Unit, *args)

pint._unpickle = _unpickle
pint._unpickle_quantity = _unpickle_quantity
pint._unpickle_unit = _unpickle_unit

zobac avatar Jul 17 '23 14:07 zobac