Search code examples
pythonpython-xarraypython-class

combine multiple xarray accessors into single one


Say I have multiple xarray accessors registered:

import xarray as xr

@xr.register_dataarray_accessor("a")
class A:
    def __init__(self, xarray_obj):
        self._obj = xarray_obj

    def print(self):
        return "A"

@xr.register_dataarray_accessor("b")
class B:
    def __init__(self, xarray_obj):
        self._obj = xarray_obj

    def print(self):
        return "b"
        
@xr.register_dataarray_accessor("c")
class C:
    def __init__(self, xarray_obj):
        self._obj = xarray_obj

    def print(self):
        return "c"

xr.DataArray().a.print()
# a

xr.DataArray().b.print()
# b

xr.DataArray().c.print()
# c

is there a good way to "move" them into a single accessor that acts as a sort of parent accessor?

Assuming this parent would be called z, my desired result would look like this:

xr.DataArray().z.a.print()
# a

xr.DataArray().z.b.print()
# b

xr.DataArray().z.c.print()
# c

Edit:

This is just a toy example, of course all the nested objects need to have access to the self._obj that Z gets.


Solution

  • The objects you register as accessors can do whatever you'd like, including having their own methods, attributes, etc. Some entire packages are available as xarray accessors - I'm thinking of rioxarray, where da.rio has a huge number of features.

    So in your case, you could simply do the following:

    class Printable:
        def __init__(self, obj):
            self._obj = obj
    
        def print_me(self):
            print(
              type(self).__name__.lower()
              + '\n'
               + str(self._obj)
            )
    
    class A(Printable):
        pass
    
    class B(Printable):
        pass
    
    class C(Printable):
        pass
    
    @xr.register_dataarray_accessor("z")
    class Z:
        def __init__(self, xarray_obj):
            self._obj = xarray_obj
    
            self.a = A(xarray_obj)
            self.b = B(xarray_obj)
            self.c = C(xarray_obj)
    

    This has the behavior you were looking for:

    In [6]: xr.DataArray().z.a.print_me()
    a
    <xarray.DataArray ()>
    array(nan)
    
    In [7]: xr.DataArray().z.b.print_me()
    b
    <xarray.DataArray ()>
    array(nan)
    
    In [8]: xr.DataArray().z.c.print_me()
    c
    <xarray.DataArray ()>
    array(nan)
    

    Have to call out that I'd strongly recommend against using reserved keywords as method names. I know this is a dummy example but I've changed these to print_me just to play it safe :)