combine multiple xarray accessors into single one

user_5997555 0
Viewed 50

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.

python-xarray python-class python
1 Answers

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 :)

user_38887190
Related