Merge pull request #15 from eyaltrabelsi/master

support xarray

Author: Zsailer

pandas_flavor/__init__.py

@@ -2,3 +2,6 @@
                        register_series_accessor,
                        register_dataframe_method,
                        register_dataframe_accessor)
+from .xarray import (register_xarray_dataarray_method, 
+                     register_xarray_dataset_method)
+

pandas_flavor/xarray.py

@@ -0,0 +1,39 @@
+
+import xarray as xr
+
+from xarray import register_dataarray_accessor, register_dataset_accessor
+from functools import wraps
+
+
+def make_accessor_wrapper(method):
+    """
+    Makes an XArray-compatible accessor to wrap a method to be added to an
+    xr.DataArray, xr.Dataset, or both.
+    :param method: A method which takes an XArray object and needed parameters.
+    :return: The result of calling ``method``.
+    """
+
+    class XRAccessor:
+        def __init__(self, xr_obj):
+            self._xr_obj = xr_obj
+
+        @wraps(method)
+        def __call__(self, *args, **kwargs):
+            return method(self._xr_obj, *args, **kwargs)
+
+    return XRAccessor
+
+
+def register_xarray_dataarray_method(method: callable):
+    accessor_wrapper = make_accessor_wrapper(method)
+    register_dataarray_accessor(method.__name__)(accessor_wrapper)
+
+    return method
+
+
+def register_xarray_dataset_method(method: callable):
+    accessor_wrapper = make_accessor_wrapper(method)
+    register_dataset_accessor(method.__name__)(accessor_wrapper)
+
+    return method
+

setup.py

@@ -19,7 +19,7 @@
 AUTHOR = 'Zach Sailer'
 
 # What packages are required for this module to be executed?
-REQUIRED = ["pandas"]
+REQUIRED = ["pandas","xarray"]
 
 # The rest you shouldn't have to touch too much :)
 # ------------------------------------------------
@@ -82,12 +82,6 @@ def run(self):
     author_email=EMAIL,
     url=URL,
     packages=find_packages(exclude=('tests',)),
-    # If your package is a single module, use this instead of 'packages':
-    # py_modules=['mypackage'],
-
-    # entry_points={
-    #     'console_scripts': ['mycli=mymodule:cli'],
-    # },
     install_requires=REQUIRED,
     include_package_data=True,
     license='MIT',