Dramatically speed up YAML parsing for setenv.sh

If available, this change makes YAML parsing use C bindings from LibYAML.
This is strongly recommended by the PyYAML team, since LibYAML bindings
are "much faster than the pure Python version."

This change will speed up the `./setenv.sh` script considerably: on my
system, `mk_shell_variables.py` decreases runtime from 450ms down to
70ms. That script takes the overwhelming majority of execution time for
`setenv.sh`.

On Ubuntu systems, `LibYAML` can be used by installing from apt:

    apt install libyaml-dev

Homebrew users can run:

    brew install libyaml

After doing so, `pyyaml` should be re-installed.

This change preserves the safe loading of files, which should *always*
be done when loading any untrusted YAML (and is good practice even if
the source is trusted). The change also doesn't *require* that one have
the C bindings.

Author: DavidCain

clrenv/__init__.py

@@ -1,12 +1,13 @@
 import yaml
 
 from .lazy_env import get_env, LazyEnv
+from .load import safe_load
 from .path import find_environment_path
 
 
 def mapping():
     with open(find_environment_path()) as f:
-        return yaml.safe_load(f.read())['mapping']
+        return safe_load(f.read())['mapping']
 
 env = LazyEnv()
 get_env = get_env

clrenv/lazy_env.py

@@ -12,8 +12,8 @@
 from botocore.exceptions import EndpointConnectionError
 
 from munch import Munch, munchify
-import yaml
 
+from .load import safe_load
 from .path import find_environment_path, find_user_environment_paths
 from functools import reduce
 
@@ -59,7 +59,7 @@ def get_env(*mode):
     if not mode in _env:
         y = (_load_current_environment(),)
         upaths = find_user_environment_paths()
-        y = tuple(yaml.safe_load(open(p).read()) for p in upaths if os.path.isfile(p)) + y
+        y = tuple(safe_load(open(p).read()) for p in upaths if os.path.isfile(p)) + y
 
         assignments = tuple(m for m in mode if m.find('=') != -1)
         mode = tuple(m for m in mode if m.find('=') == -1)
@@ -78,7 +78,7 @@ def get_env(*mode):
         e = _merged(*dicts)
 
         for k, v in overrides:
-            for pytype in (yaml.safe_load, eval, int, float, str):
+            for pytype in (safe_load, eval, int, float, str):
                 try:
                     pyval = pytype(v)
                     break
@@ -146,7 +146,7 @@ def _setattr_rec(d, k, v):
 
 def _load_current_environment():
     with open(find_environment_path()) as f:
-        environment = yaml.safe_load(f.read())
+        environment = safe_load(f.read())
     return environment
 
 _kf_dict_cache = {}

clrenv/load.py

@@ -0,0 +1,21 @@
+"""
+This module exists to provide a faster, yet equally secure version of `safe_load`
+"""
+import yaml
+
+try:
+    # If available, use the C bindings for far, far faster loading
+    # See: https://pyyaml.org/wiki/PyYAMLDocumentation
+    from yaml import CSafeLoader as SafeLoader
+except ImportError:
+    # If the C bindings aren't available, fall back to the "much slower" Python bindings
+    from yaml import SafeLoader
+
+
+def safe_load(str_content):
+    """ Safely load YAML, doing so quickly with C bindings if available.
+
+    By default, `yaml.safe_load()` uses the (slower) Python bindings.
+    This method is a stand-in replacement that can be considerably faster.
+    """
+    return yaml.load(str_content, Loader=SafeLoader)

setup.py

@@ -7,7 +7,7 @@
 
 
 setup(name = "clrenv",
-      version = "0.1.6",
+      version = "0.1.7",
       description = "A tool to give easy access to environment yaml file to python.",
       author = "Color Genomics",
       author_email = "[email protected]",