Commits

Arne Babenhauserheide  committed fdcbd5d

Added get_all_datasets to get all stored values in the internal format (sorted by storage file).

  • Participants
  • Parent commits 51bf76b

Comments (0)

Files changed (1)

 
 import yaml
 from os.path import join
-from os import makedirs
+from os import makedirs, listdir
 
 def key_to_paths(key,
                 sites_dir="sites",
         raise IOError("Can’t write to the states datafile. Please check if the states folder exists:" + states_dir)
 
 
+def get_all_datasets(sites_dir="sites",
+                     states_dir="states"):
+    """Store the data for a given key.
+
+    >>> _test_clean_dirs()
+    >>> _test_create_test_dirs()
+    >>> get_all_datasets(sites_dir="test_sites", states_dir="test_states")
+    ({}, {})
+    >>> sites = [{1: 2}, {5:6}, {'path': 'yum'}]
+    >>> states = [{5:7}]
+    >>> put_dataset_for_key("USK@123/foo/bar", sites, states, sites_dir="test_sites", states_dir="test_states")
+    >>> si, st = get_all_datasets(sites_dir="test_sites", states_dir="test_states")
+    >>> si.keys(), st.keys()
+    (dict_keys(['USK@123']), dict_keys(['USK@123.state']))
+    """
+    datasets_sites = {}
+    datasets_states = {}
+    for i in listdir(sites_dir):
+        sites_path = join(sites_dir, i)
+        # sites data
+        try: 
+            with open(sites_path) as f: 
+                yaml_data = f.read()
+            sites = yaml.safe_load_all(yaml_data)
+        except IOError:
+            sites = (i for i in []) # empty generator
+        datasets_sites[i] = sites
+
+    for i in listdir(states_dir):
+        states_path = join(states_dir, i)
+        # states data
+        try: 
+            with open(states_path) as f: 
+                yaml_data = f.read()
+            states = yaml.safe_load_all(yaml_data)
+        except IOError:
+            states = (i for i in []) # empty generator
+        datasets_states[i] = states
+
+    return datasets_sites, datasets_states
+
+
 def _test_clean_dirs():
     """Cleanup the test folders before the tests."""
     from subprocess import call