From cdb8af6cb68bf04c5879cb6fa501a76aa96b9b8e Mon Sep 17 00:00:00 2001
From: Nathan Collier <nathaniel.collier@gmail.com>
Date: Mon, 22 Jan 2024 17:33:01 -0500
Subject: [PATCH] conditional import of tqdm_notebook (#13)

* conditional import of tqdm_notebook

* added missing widget for notebooks

* fix max filename length in downloads
---
 ci/environment-docs.yml  |  1 +
 intake_esgf/__init__.py  | 22 +++++++++++++++++++---
 intake_esgf/base.py      | 15 +++++++++++++--
 intake_esgf/catalog.py   |  7 ++++++-
 intake_esgf/operators.py |  7 ++++++-
 5 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/ci/environment-docs.yml b/ci/environment-docs.yml
index 85d93d3..ed8a197 100644
--- a/ci/environment-docs.yml
+++ b/ci/environment-docs.yml
@@ -11,6 +11,7 @@ dependencies:
   - distributed
   - netCDF4
   - matplotlib
+  - ipywidgets
   - pip
   - pip:
       - myst_nb
diff --git a/intake_esgf/__init__.py b/intake_esgf/__init__.py
index eee6be8..4921ce3 100644
--- a/intake_esgf/__init__.py
+++ b/intake_esgf/__init__.py
@@ -4,9 +4,25 @@
 
 import xarray as xr
 
-from intake_esgf.catalog import ESGFCatalog
-
 warnings.simplefilter("ignore", category=xr.SerializationWarning)
 
 
-__all__ = ["ESGFCatalog"]
+def in_notebook() -> bool:
+    """Check if the code is running in a jupyter notebook"""
+    try:
+        shell = get_ipython().__class__.__name__
+        if shell == "ZMQInteractiveShell":  # Jupyter notebook, Spyder or qtconsole
+            return True
+        elif shell == "TerminalInteractiveShell":
+            return False  # Terminal running IPython
+        else:
+            return False  # Other type (?)
+    except NameError:
+        return False  # Probably standard Python interpreter
+
+
+IN_NOTEBOOK = in_notebook()
+
+from intake_esgf.catalog import ESGFCatalog  # noqa
+
+__all__ = ["ESGFCatalog", "IN_NOTEBOOK"]
diff --git a/intake_esgf/base.py b/intake_esgf/base.py
index 2ea0255..cc19830 100644
--- a/intake_esgf/base.py
+++ b/intake_esgf/base.py
@@ -10,8 +10,8 @@
 import pandas as pd
 import requests
 import xarray as xr
-from tqdm import tqdm
 
+from intake_esgf import IN_NOTEBOOK
 from intake_esgf.database import (
     get_download_rate_dataframe,
     log_download_information,
@@ -20,6 +20,11 @@
 from intake_esgf.exceptions import NoSearchResults
 from intake_esgf.logging import setup_logging
 
+if IN_NOTEBOOK:
+    from tqdm import tqdm_notebook as tqdm
+else:
+    from tqdm import tqdm
+
 bar_format = "{desc:>20}: {percentage:3.0f}%|{bar}|{n_fmt}/{total_fmt} [{rate_fmt:>15s}{postfix}]"
 
 
@@ -93,6 +98,12 @@ def download_and_verify(
     """Download the url to a local file and check for validity, removing if not."""
     if not isinstance(local_file, Path):
         local_file = Path(local_file)
+    max_file_length = 40
+    desc = (
+        local_file.name
+        if len(local_file.name) < max_file_length
+        else f"{local_file.name[:(max_file_length-3)]}..."
+    )
     local_file.parent.mkdir(parents=True, exist_ok=True)
     resp = requests.get(url, stream=True, timeout=10)
     resp.raise_for_status()
@@ -104,7 +115,7 @@ def download_and_verify(
             total=content_length,
             unit="B",
             unit_scale=True,
-            desc=local_file.name,
+            desc=desc,
             ascii=False,
         ) as pbar:
             for chunk in resp.iter_content(chunk_size=1024):
diff --git a/intake_esgf/catalog.py b/intake_esgf/catalog.py
index c30641c..0c13195 100644
--- a/intake_esgf/catalog.py
+++ b/intake_esgf/catalog.py
@@ -11,8 +11,8 @@
 import requests
 import xarray as xr
 from datatree import DataTree
-from tqdm import tqdm
 
+from intake_esgf import IN_NOTEBOOK
 from intake_esgf.base import (
     add_cell_measures,
     bar_format,
@@ -24,6 +24,11 @@
 from intake_esgf.database import create_download_database, get_download_rate_dataframe
 from intake_esgf.logging import setup_logging
 
+if IN_NOTEBOOK:
+    from tqdm import tqdm_notebook as tqdm
+else:
+    from tqdm import tqdm
+
 
 class ESGFCatalog:
     """A data catalog for searching ESGF nodes and downloading data.
diff --git a/intake_esgf/operators.py b/intake_esgf/operators.py
index 29a789f..2d75b9f 100644
--- a/intake_esgf/operators.py
+++ b/intake_esgf/operators.py
@@ -3,10 +3,15 @@
 
 import pandas as pd
 import xarray as xr
-from tqdm import tqdm
 
+from intake_esgf import IN_NOTEBOOK
 from intake_esgf.base import bar_format, get_cell_measure, get_search_criteria
 
+if IN_NOTEBOOK:
+    from tqdm import tqdm_notebook as tqdm
+else:
+    from tqdm import tqdm
+
 
 def global_sum(
     dsd: Union[dict[str, xr.Dataset], xr.Dataset], quiet: bool = False