diff --git a/binarycpython/utils/functions.py b/binarycpython/utils/functions.py index b46b12ec177baa43d585030a8e1fbd2a08752811..619354e78132eb648e7eb4a5145a259faebfaacd 100644 --- a/binarycpython/utils/functions.py +++ b/binarycpython/utils/functions.py @@ -18,9 +18,10 @@ import copy import datetime as dt import gc import gzip -import inspect +from halo import Halo import h5py import humanize +import inspect from io import StringIO import json import numpy as np @@ -33,6 +34,7 @@ import sys import subprocess import tempfile import time +from tqdm import tqdm import types from typing import Union, Any @@ -2253,10 +2255,14 @@ class BinaryCEncoder(json.JSONEncoder): # Let the base class default method raise the TypeError return json.JSONEncoder.default(self, o) -def load_ensemble(filename,convert_float_keys=True): +def load_ensemble(filename,convert_float_keys=True,select_keys=None): """ Function to load an ensemeble file, even if it is compressed, and return its contents to as a Python dictionary. + + Args: + convert_float_keys : if True, converts strings to floats. + select_keys : a list of keys to be selected from the ensemble. """ if(filename.endswith('.bz2')): jfile = bz2.open(filename,'rt') @@ -2264,7 +2270,30 @@ def load_ensemble(filename,convert_float_keys=True): jfile = gzip.open(filename,'rt') else: jfile = open(filename,'rt') - data = json.load(jfile) + + + + # load with some info to the terminal + print("Loading JSON...") + _loaded = False + def _hook(obj): + nonlocal _loaded + if _loaded == False: + _loaded = True + print("\nLoaded JSON data, now putting in a dictionary") + return obj + with Halo(text='Loading', interval=250, spinner='moon',color='yellow'): + data = json.load(jfile, + object_hook=_hook) + + # strip non-selected keys, if a list is given in select_keys + if select_keys: + keys = list(data['ensemble'].keys()) + for key in keys: + if not key in select_keys: + del data['ensemble'][key] + + # perhaps convert floats? if convert_float_keys == False: return data else: diff --git a/requirements.txt b/requirements.txt index d2cfebe3bc46d8b2edeef17e5928260df2c87c3a..bbae3d56309cdf96ff3610682d199e4151fb9cd1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,6 +10,7 @@ cycler==0.10.0 decorator==4.4.1 dill==0.3.1.1 docutils==0.15.2 +halo==0.0.31 h5py==2.10.0 hawkmoth==0.4 humanize==3.12.0 diff --git a/setup.py b/setup.py index f33a42870c18885b2d432aa152799aa4abfad2c4..5257ba5404b45305d5629423f880b2ed849788bd 100644 --- a/setup.py +++ b/setup.py @@ -262,7 +262,8 @@ setup( "psutil", "colorama", "strip-ansi", - "humanize" + "humanize", + "halo" ], include_package_data=True, ext_modules=[BINARY_C_PYTHON_API_MODULE], # binary_c must be loaded