Add tests

Author: andersy005

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "python.pythonPath": "/Users/abanihi/opt/miniconda3/envs/pangeo/bin/python"
+}

intake_cmip5/__init__.py

@@ -1,8 +1,10 @@
 #!/usr/bin/env python
 """Top-level package for intake_cmip5."""
 from ._version import get_versions
+import intake_cmip5
+from intake_cmip5 import generate_database
 
 __version__ = get_versions()["version"]
 del get_versions
 
-__all__ = []
+__all__ = ["intake_cmip5", "generate_database"]

intake_cmip5/_version.py

@@ -68,8 +68,7 @@ def decorate(f):
     return decorate
 
 
-def run_command(commands, args, cwd=None, verbose=False,
-                hide_stderr=False, env=None):
+def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, env=None):
     """Call the given command(s)."""
     assert isinstance(commands, list)
     p = None
@@ -121,7 +120,7 @@ def versions_from_parentdir(parentdir_prefix, root, verbose):
         dirname = os.path.basename(root)
         if dirname.startswith(parentdir_prefix):
             return {
-                "version": dirname[len(parentdir_prefix):],
+                "version": dirname[len(parentdir_prefix) :],
                 "full-revisionid": None,
                 "dirty": False,
                 "error": None,
@@ -191,7 +190,7 @@ def git_versions_from_keywords(keywords, tag_prefix, verbose):
     # starting in git-1.8.3, tags are listed as "tag: foo-1.0" instead of
     # just "foo-1.0". If we see a "tag: " prefix, prefer those.
     TAG = "tag: "
-    tags = set([r[len(TAG):] for r in refs if r.startswith(TAG)])
+    tags = set([r[len(TAG) :] for r in refs if r.startswith(TAG)])
     if not tags:
         # Either we're using git < 1.8.3, or there really are no tags. We use
         # a heuristic: assume all version tags have a digit. The old git %d
@@ -208,7 +207,7 @@ def git_versions_from_keywords(keywords, tag_prefix, verbose):
     for ref in sorted(tags):
         # sorting will prefer e.g. "2.0" over "2.0rc1"
         if ref.startswith(tag_prefix):
-            r = ref[len(tag_prefix):]
+            r = ref[len(tag_prefix) :]
             if verbose:
                 print("picking %s" % r)
             return {
@@ -242,8 +241,7 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, run_command=run_command):
     if sys.platform == "win32":
         GITS = ["git.cmd", "git.exe"]
 
-    out, rc = run_command(
-        GITS, ["rev-parse", "--git-dir"], cwd=root, hide_stderr=True)
+    out, rc = run_command(GITS, ["rev-parse", "--git-dir"], cwd=root, hide_stderr=True)
     if rc != 0:
         if verbose:
             print("Directory %s not under git control" % root)
@@ -313,7 +311,7 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, run_command=run_command):
                 tag_prefix,
             )
             return pieces
-        pieces["closest-tag"] = full_tag[len(tag_prefix):]
+        pieces["closest-tag"] = full_tag[len(tag_prefix) :]
 
         # distance: number of commits since tag
         pieces["distance"] = int(mo.group(2))
@@ -324,8 +322,7 @@ def git_pieces_from_vcs(tag_prefix, root, verbose, run_command=run_command):
     else:
         # HEX: no tags
         pieces["closest-tag"] = None
-        count_out, rc = run_command(
-            GITS, ["rev-list", "HEAD", "--count"], cwd=root)
+        count_out, rc = run_command(GITS, ["rev-list", "HEAD", "--count"], cwd=root)
         pieces["distance"] = int(count_out)  # total number of commits
 
     # commit date: see ISO-8601 comment in git_versions_from_keywords()
@@ -522,8 +519,7 @@ def get_versions():
     verbose = cfg.verbose
 
     try:
-        return git_versions_from_keywords(
-            get_keywords(), cfg.tag_prefix, verbose)
+        return git_versions_from_keywords(get_keywords(), cfg.tag_prefix, verbose)
     except NotThisMethod:
         pass
 
@@ -561,4 +557,4 @@ def get_versions():
         "dirty": None,
         "error": "unable to compute version",
         "date": None,
-    }
+    }

intake_cmip5/generate_database.py

@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-import os 
+import os
 import pandas as pd
 from dask import delayed
 import re
@@ -9,55 +9,77 @@
 import shutil
 
 
-
 HOME = os.environ["HOME"]
 INTAKE_CMIP5_DIR = f"{HOME}/.intake_cmip5"
 
+
 @functools.lru_cache(maxsize=1024, typed=False)
 def _parse_dirs(root_dir):
-    institution_dirs = [os.path.join(root_dir, activity, institution)
-                           for activity in os.listdir(root_dir)
-                           for institution in os.listdir(os.path.join(root_dir, activity))
-                           if os.path.isdir(os.path.join(root_dir, activity, institution))]
-        
-    model_dirs = [os.path.join(institution_dir, model)
-                      for institution_dir in institution_dirs
-                      for model in os.listdir(institution_dir)
-                      if os.path.isdir(os.path.join(institution_dir, model))]
-        
-    experiment_dirs = [os.path.join(model_dir, exp)
-                           for model_dir in model_dirs
-                           for exp in os.listdir(model_dir)
-                           if os.path.isdir(os.path.join(model_dir, exp))]
-        
-    freq_dirs = [os.path.join(experiment_dir, freq)
-                     for experiment_dir in experiment_dirs
-                     for freq in os.listdir(experiment_dir)
-                     if os.path.isdir(os.path.join(experiment_dir, freq))]
-        
-    realm_dirs = [os.path.join(freq_dir, realm)
-                      for freq_dir in freq_dirs
-                      for realm in os.listdir(freq_dir)
-                      if os.path.isdir(os.path.join(freq_dir, realm))]
-        
+    institution_dirs = [
+        os.path.join(root_dir, activity, institution)
+        for activity in os.listdir(root_dir)
+        for institution in os.listdir(os.path.join(root_dir, activity))
+        if os.path.isdir(os.path.join(root_dir, activity, institution))
+    ]
+
+    model_dirs = [
+        os.path.join(institution_dir, model)
+        for institution_dir in institution_dirs
+        for model in os.listdir(institution_dir)
+        if os.path.isdir(os.path.join(institution_dir, model))
+    ]
+
+    experiment_dirs = [
+        os.path.join(model_dir, exp)
+        for model_dir in model_dirs
+        for exp in os.listdir(model_dir)
+        if os.path.isdir(os.path.join(model_dir, exp))
+    ]
+
+    freq_dirs = [
+        os.path.join(experiment_dir, freq)
+        for experiment_dir in experiment_dirs
+        for freq in os.listdir(experiment_dir)
+        if os.path.isdir(os.path.join(experiment_dir, freq))
+    ]
+
+    realm_dirs = [
+        os.path.join(freq_dir, realm)
+        for freq_dir in freq_dirs
+        for realm in os.listdir(freq_dir)
+        if os.path.isdir(os.path.join(freq_dir, realm))
+    ]
+
     return realm_dirs
-    
+
+
 def _get_entry(directory):
-    dir_split = directory.split('/')
+    dir_split = directory.split("/")
     entry = {}
-    entry['realm'] = dir_split[-1]
-    entry['frequency'] = dir_split[-2]
-    entry['experiment'] = dir_split[-3]
-    entry['model'] = dir_split[-4]
-    entry['institution'] = dir_split[-5]
+    entry["realm"] = dir_split[-1]
+    entry["frequency"] = dir_split[-2]
+    entry["experiment"] = dir_split[-3]
+    entry["model"] = dir_split[-4]
+    entry["institution"] = dir_split[-5]
     return entry
-        
+
+
 @delayed
 def parse_directory(directory):
-    exclude = set(["files", "latests"]) # directories to exclude
+    exclude = set(["files", "latests"])  # directories to exclude
 
-    columns = ["ensemble", "experiment", "file_basename", "file_fullpath", 
-                      "frequency", "institution", "model", "root", "realm", "varname"]
+    columns = [
+        "ensemble",
+        "experiment",
+        "file_basename",
+        "file_fullpath",
+        "frequency",
+        "institution",
+        "model",
+        "root",
+        "realm",
+        "varname",
+    ]
     df = pd.DataFrame(columns=columns)
 
     entry = _get_entry(directory)
@@ -68,17 +90,18 @@ def parse_directory(directory):
         if not files:
             continue
         sfiles = sorted([f for f in files if os.path.splitext(f)[1] == ".nc"])
-        if not sfiles: continue
+        if not sfiles:
+            continue
 
         fs = []
         for f in sfiles:
             try:
                 f_split = f.split("_")
-                entry['varname'] = f_split[0]
-                entry['ensemble'] = f_split[-2]
-                entry['root'] = root
-                entry['file_basename'] = f
-                entry['file_fullpath'] = os.path.join(root, f)
+                entry["varname"] = f_split[0]
+                entry["ensemble"] = f_split[-2]
+                entry["root"] = root
+                entry["file_basename"] = f
+                entry["file_fullpath"] = os.path.join(root, f)
                 fs.append(entry)
             except:
                 continue
@@ -88,34 +111,45 @@ def parse_directory(directory):
         else:
             temp_df = pd.DataFrame()
             temp_df.columns = df.columns
-        df = pd.concat([temp_df, df], ignore_index=True)
+        df = pd.concat([temp_df, df], ignore_index=True, sort=False)
     return df
 
-def _persist_database(df):
-    vYYYYMMDD = r'v\d{4}\d{2}\d{2}'
-    vN = r'v\d{1}'
-    v = re.compile( "|".join([vYYYYMMDD, vN])) # Combine both regex into one
+
+def _persist_database(df, path):
+    vYYYYMMDD = (
+        r"v\d{4}\d{2}\d{2}"
+    )  # TODO: Very dangerous in case the root dir matches the pattern
+    vN = r"v\d{1}"
+    v = re.compile("|".join([vYYYYMMDD, vN]))  # Combine both regex into one
     df["version"] = df.root.str.findall(v)
-    df["version"] = df["version"].apply(lambda x: x[0] if x else 'v0')
-    sorted_df = df.sort_values("version").drop_duplicates(subset="file_basename", keep="last")\
-                  .reset_index(drop=True)
+    df["version"] = df["version"].apply(lambda x: x[0] if x else "v0")
+    sorted_df = (
+        df.sort_values("version")
+        .drop_duplicates(subset="file_basename", keep="last")
+        .reset_index(drop=True)
+    )
+
+    if path:
+        INTAKE_CMIP5_DIR = path
+
     print(f"**** Persisting CMIP5 database in {INTAKE_CMIP5_DIR} ****")
 
     if os.path.isdir(INTAKE_CMIP5_DIR):
         shutil.rmtree(INTAKE_CMIP5_DIR)
     os.makedirs(INTAKE_CMIP5_DIR, exist_ok=True)
-    
+
     sorted_df.to_csv(f"{INTAKE_CMIP5_DIR}/clean_cmip5_database.csv", index=False)
     df.to_csv(f"{INTAKE_CMIP5_DIR}/raw_cmip5_database.csv", index=False)
-    
+
     return sorted_df
 
-def create_CMIP5Database(root_dir=None):
+
+def create_CMIP5Database(root_dir=None, db_path=None):
     if not os.path.exists(root_dir):
         raise NotADirectoryError(f"{root_dir} does not exist")
-        
+
     dirs = _parse_dirs(root_dir)
     dfs = [parse_directory(directory) for directory in dirs]
     df = dd.from_delayed(dfs).compute()
-    df = _persist_database(df)
-    return df
+    df = _persist_database(df, db_path)
+    return df

setup.cfg

@@ -15,7 +15,11 @@ replace = __version__ = '{new_version}'
 universal = 1
 
 [flake8]
-exclude = docs
+exclude = docs,versioneer.py,intake_cmip5/._version.py
+ignore = E203, E266, E501, W503, F401, E722
+max-line-length = 88
+max-complexity = 18
+select = B,C,E,F,W,T4,B9
 
 [aliases]
 # Define setup.py command aliases here

setup.py

@@ -9,26 +9,26 @@
 with open("README.md", encoding="utf-8") as readme_file:
     readme = readme_file.read()
 
-requirements = open('requirements.txt').read().strip().split('\n')
+requirements = open("requirements.txt").read().strip().split("\n")
 
 setup(
     maintainer="Anderson Banihirwe",
     maintainer_email="[email protected]",
-    description='An intake plugin for loading CMIP5 data sets',
+    description="An intake plugin for loading CMIP5 data sets",
     install_requires=requirements,
     license="https://github.com/NCAR/intake-cmip5/blob/master/LICENSE.rst",
     long_description=readme,
     long_description_content_type="text/markdown",
     keywords=["cmip5", "intake"],
     name="intake-cmip5",
     packages=find_packages(),
-    py_modules=['intake_cmip5'],
-    package_data={'': ['*.yml', '*.yaml']},
+    py_modules=["intake_cmip5"],
+    package_data={"": ["*.yml", "*.yaml", "*.csv"]},
     include_package_data=True,
     url="https://github.com/NCAR/intake-cmip5",
     version=versioneer.get_version(),
     cmdclass=versioneer.get_cmdclass(),
     entry_points="""
       """,
     zip_safe=False,
-)
+)