#200: added support for xarray datasets

dtale/cli/clickutils.py

@@ -35,10 +35,6 @@ def setup_logging(logfile, log_level, verbose=False):
     else:
         log_level = LOG_LEVELS['info']
 
-    # for pkg in ['_plotly_utils', 'asyncio', 'concurrent', 'matplotlib', 'parso', 'past', 'prompt_toolkit', 'requests',
-    #             'tornado', 'urllib3']:
-    #     logging.getLogger(pkg).propagate = True
-
     logging.getLogger().handlers = []
 
     fmt = "%(asctime)s - %(levelname)-8s - %(message)s"

dtale/global_state.py

@@ -2,6 +2,8 @@
 from six.moves.collections_abc import MutableMapping
 
 DATA = {}
+DATASETS = {}
+DATASET_DIM = {}
 DTYPES = {}
 SETTINGS = {}
 METADATA = {}
@@ -17,6 +19,22 @@ def get_data(data_id=None):
     return DATA.get(data_id)
 
 
+def get_dataset(data_id=None):
+    global DATASETS
+
+    if data_id is None:
+        return _as_dict(DATASETS)
+    return DATASETS.get(data_id)
+
+
+def get_dataset_dim(data_id=None):
+    global DATASET_DIM
+
+    if data_id is None:
+        return _as_dict(DATASET_DIM)
+    return DATASET_DIM.get(data_id)
+
+
 def get_dtypes(data_id=None):
     global DTYPES
 
@@ -63,6 +81,18 @@ def set_data(data_id, val):
     DATA[data_id] = val
 
 
+def set_dataset(data_id, val):
+    global DATASETS
+
+    DATASETS[data_id] = val
+
+
+def set_dataset_dim(data_id, val):
+    global DATASET_DIM
+
+    DATASET_DIM[data_id] = val
+
+
 def set_dtypes(data_id, val):
     global DTYPES
 
@@ -100,17 +130,19 @@ def cleanup(data_id=None):
     :param port: integer string for a D-Tale process's port
     :type port: str
     """
-    global DATA, DTYPES, SETTINGS, METADATA, CONTEXT_VARIABLES, HISTORY
+    global DATA, DATASETS, DATASET_DIM, DTYPES, SETTINGS, METADATA, CONTEXT_VARIABLES, HISTORY
 
     if data_id is None:
         DATA.clear()
+        DATASETS.clear()
+        DATASET_DIM.clear()
         SETTINGS.clear()
         DTYPES.clear()
         METADATA.clear()
         CONTEXT_VARIABLES.clear()
         HISTORY.clear()
     else:
-        for store in [DATA, DTYPES, SETTINGS, METADATA, CONTEXT_VARIABLES, HISTORY]:
+        for store in [DATA, DATASETS, DATASET_DIM, DTYPES, SETTINGS, METADATA, CONTEXT_VARIABLES, HISTORY]:
             if data_id in store:
                 del store[data_id]
 

dtale/templates/dtale/base.html

@@ -30,6 +30,8 @@
         </span>
         {% endif %}
         <input type="hidden" id="data_id" value="{{data_id}}" />
+        <input type="hidden" id="xarray" value="{{xarray}}" />
+        <input type="hidden" id="xarray_dim" value="{{xarray_dim}}" />
         <input type="hidden" id="settings" value="{{settings}}" />
         <input type="hidden" id="version" value="{{version}}" />
         <input type="hidden" id="hide_shutdown" value="{{config.HIDE_SHUTDOWN}}" />

dtale/utils.py

@@ -729,18 +729,33 @@ def build_code_export(data_id, imports='import pandas as pd\n\n', query=None):
     settings = global_state.get_settings(data_id) or {}
     ctxt_vars = global_state.get_context_variables(data_id)
 
-    startup_code = settings.get('startup_code')
-    startup_code = '# Data Re-shaping\n{}\n\n'.format(startup_code) if startup_code else ''
+    startup_code = settings.get('startup_code') or ''
+    xarray_setup = ''
+    if data_id in global_state.DATASETS:
+        xarray_dims = global_state.get_dataset_dim(data_id)
+        if len(xarray_dims):
+            xarray_setup = (
+                'df = ds.sel({selectors}).to_dataframe()\n'
+                "df = df.reset_index().drop('index', axis=1, errors='ignore')\n"
+                'df = df.set_index(list(ds.dims.keys()))\n'
+            ).format(selectors=', '.join("{}='{}'".format(k, v) for k, v in xarray_dims.items()))
+        else:
+            xarray_setup = (
+                'df = ds.to_dataframe()\n'
+                "df = df.reset_index().drop('index', axis=1, errors='ignore')\n"
+                'df = df.set_index(list(ds.dims.keys()))\n'
+            )
     startup_str = (
         "# DISCLAIMER: 'df' refers to the data you passed in when calling 'dtale.show'\n\n"
         '{imports}'
+        '{xarray_setup}'
         '{startup}'
         'if isinstance(df, (pd.DatetimeIndex, pd.MultiIndex)):\n'
         '\tdf = df.to_frame(index=False)\n\n'
         '# remove any pre-existing indices for ease of use in the D-Tale code, but this is not required\n'
         "df = df.reset_index().drop('index', axis=1, errors='ignore')\n"
         'df.columns = [str(c) for c in df.columns]  # update columns to strings in case they are numbers\n'
-    ).format(imports=imports, startup=startup_code)
+    ).format(imports=imports, xarray_setup=xarray_setup, startup=startup_code)
     final_history = [startup_str] + history
     final_query = query
     if final_query is None: