-
Notifications
You must be signed in to change notification settings - Fork 0
/
datapipes.py
343 lines (292 loc) · 11.3 KB
/
datapipes.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
#!/usr/bin/env python3
import os
import re
import h5py
import json
import inspect
import subprocess
import numpy as np
from functools import wraps
from printech import *
def load_from_json(filename):
'''Load data from a json file into a dictionary
omitting lines that start with //. This is a simple
loader that does not support nested dictionaries.
Parameters
----------
filename: str
The path to the json file to load.
Returns
-------
the_dictionary: dict
The dictionary with the contents of the json file.
'''
lines = []
with open(filename, 'r') as file:
for line in file:
clear_line = line.split('//')[0].strip()
if clear_line:
lines.append(clear_line)
clear_lines = '\n'.join(lines)
the_dictionary = json.loads(clear_lines)
return the_dictionary
class CustomJSONEncoder(json.JSONEncoder):
def default(self, obj):
if isinstance(obj, np.ndarray):
# Convert numpy arrays to lists
return obj.tolist()
# Let the base class default method raise the TypeError
return json.JSONEncoder.default(self, obj)
def save_to_json(filename, source_dict, header=''):
'''
This function saves a dictionary to a JSON file with the given filename.
It admits values to be int, float, str, list, dict, or numpy arrays.
This is a simple wrapper around json.dump that uses a custom encoder that
accounts for numpy arrays.
Parameters
----------
filename: str
The path to the JSON file to save.
source_dict: dict
The dictionary to save to the JSON file.
header: str
A string to be saved as a comment in the first line of the file.
Returns
-------
None
'''
with open(filename, 'w', encoding='utf-8') as file:
json.dump(source_dict, file, ensure_ascii=False, cls=CustomJSONEncoder, indent=4)
return None
def save_to_h5(data, filename, group=None, comments='', overwrite=False):
'''
Save numeric values, numpy arrays, strings, and nested dictionaries
from the provided dictionary to an H5 file.
If a value of the provided dictionary is a list, then
the list is converted to a dictionary with keys equal
to the indices of the lists (as strings).
Parameters
----------
data: dict
The input dictionary.
filename: str
The path to the H5 file to save.
group: (h5py.Group, optional)
Current group for recursion. Defaults to None.
Returns
-------
None
Example
-------
.. code-block:: python
data_dict = {
'a': [1,2,3],
'b': {
'c': 'hello',
'd': np.array([1, 2, 3]),
'e': {
'f': 3.14,
'g': 'world'
}
},
'h': np.array([[1, 2], [3, 4]])
}
save_to_h5(data_dict, 'output.h5')
'''
if group == None and os.path.exists(filename):
if overwrite:
print("File already exists, overwriting ...")
else:
print("File already exists, doing nothing ...")
return None
# Initialize the file/group on the first call
if group is None:
with h5py.File(filename, 'w') as h5f:
if comments:
h5f.attrs['comments'] = comments.encode('utf-8')
save_to_h5(data, filename, group=h5f)
else:
for key, value in data.items():
# If the value is a dictionary, recurse into it
if isinstance(value, dict):
if isinstance(key, int):
key = str(key)
subgroup = group.create_group(key)
save_to_h5(value, filename, group=subgroup)
# Convert lists to dictionaries
elif isinstance(value, list):
list_dict = {str(i): v for i, v in enumerate(value)}
subgroup = group.create_group(key)
save_to_h5(list_dict, filename, group=subgroup)
else:
# If the value is a string, encode it to UTF-8
if isinstance(value, str):
value = value.encode('utf-8')
group.create_dataset(key, data=value)
def load_from_h5(filename, keys=None, only_keys=False, return_comments=True):
'''
Load an H5 file and return its contents or its keys.
Parameters
----------
filename: str
The path to the H5 file to load.
keys: (list, str, or None, optional)
A list of keys, a single key, or None to load all from the H5 file.
If a single key is given, then the function does not return
a dictionary with a comment but it simply returns the value of
that single key.
group: (h5py.Group, optional)
Current group for recursion. Defaults to None.
only_keys: (bool)
If set to True, returns only the keys without loading any data. Defaults to False.
Returns
-------
dict, value, list or None
Depending on the parameters, returns a dictionary, a value, a list of keys, or None.
tuple:
A tuple containing a dictionary with the loaded data and the comments string.
'''
def retrieve_group(group):
data = {}
for key in group:
if isinstance(group[key], h5py.Group):
data[key] = retrieve_group(group[key])
else:
data[key] = group[key][()]
if isinstance(data[key], bytes):
data[key] = data[key].decode('utf-8')
return data
with h5py.File(filename, 'r') as h5f:
if only_keys:
comment = h5f.attrs.get('comments')
return list(h5f.keys()), comment
comment = h5f.attrs.get('comments')
if keys:
data = {}
for key in keys:
if key in h5f:
if isinstance(h5f[key], h5py.Group):
data[key] = retrieve_group(h5f[key])
else:
data[key] = h5f[key][()]
if isinstance(data[key], bytes):
data[key] = data[key].decode('utf-8')
else:
data = retrieve_group(h5f)
if len(data) == 1:
return list(data.values())[0]
else:
if return_comments:
return data, comment
else:
return data
def run_script_and_get_json(script_path, args_dict):
'''
This function takes the path to a script with a command-line
interface, and a dictionary with values to all the
parameters that the CLI receives.
This assumes that the names of the arguments are known
beforehand so that args_dict may be well formed.
For an alternative that automatically analyzes the script to
be run, see extract_cli_arguments_and_run.
Using the given values for the arguments this function runs
the script.
The script must save its results to a .json file and the
last line that it prints must be the name of such file.
When the script has finished running, this function will
then grab the filename of that json file, read its contents
and return them as a dictionary.
Parameters
----------
script_path: str
Path to Python script with CLI and json output.
args_dict: dict
Dictionary with keys equal to the named arguments of the
CLI and values equal to the values one wishes to input.
Returns
-------
result_data: dict
Dictionary with the parsed results of the function as
saved in the json file by the Python script.
'''
# Construct the command with named arguments
command = [script_path]
for arg, value in args_dict.items():
command.append(f"--{arg}")
command.append(str(value))
# Run the command and capture the output
process = subprocess.Popen(command,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True)
stdout_lines = []
while True:
output = process.stdout.readline()
if output == '' and process.poll() is not None:
break
if output:
clean_output = output.strip()
print(clean_output)
stdout_lines.append(clean_output)
# Parse the last line of the standard output to get the JSON filename
json_output_filename = stdout_lines[-1]
# Read the result from the JSON file
with open(json_output_filename, 'r') as file:
result_data = json.load(file)
result_data['stdout'] = '\n'.join(stdout_lines)
return result_data
def extract_cli_arguments_and_run(script_path):
'''
Given the path to a python script with a command-line
interface. This function inspects the command line
interface, which must be built with argparse and it define a
function that will take in the CLI parameters as positional
arguments, feed them to the script via the CLI, and load the
output of the script.
Eventhough the CLI might allow for optional parameters the
function returned here requires that every parameter is
given.
The script needs to output its results to a json file and
the last line that it prints must contain the name of the
file to which the function saved its results to.
Parameters
----------
script_path : str
The filepath to the script to be wrapped.
Returns
-------
cli_function_wrapper : fun
A function that takes as aruments the arguments of the
CLI and which returns a dictionary with the results
parse from json output.
'''
# Regular expression to match argparse argument definitions
arg_pattern = re.compile(r'\.add_argument\([\'"]--(\w+)')
# List to store argument names
arg_names = []
# Read the script and search for argparse definitions
with open(script_path, 'r') as file:
for line in file:
match = arg_pattern.search(line)
if match:
# Extract the argument name and add it to the list
arg_names.append(match.group(1))
# Function that takes CLI arguments, wraps them into a dictionary,
# and feeds them to run_script_and_get_json
def cli_function_template(*args):
# Construct the args dictionary from positional arguments
args_dict = dict(zip(arg_names, args))
# Call the run_script_and_get_json function with the script path and args dictionary
return run_script_and_get_json(script_path, args_dict)
# Set the __name__ attribute of the function to a generic name like 'cli_function'
cli_function_template.__name__ = 'cli_function'
# Use functools.wraps to preserve the metadata of the template function
@wraps(cli_function_template)
def cli_function_wrapper(*args):
return cli_function_template(*args)
# Update the signature of the wrapper function to match the argument names
cli_function_wrapper.__signature__ = inspect.signature(cli_function_template).replace(parameters=[
inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD)
for name in arg_names
])
return cli_function_wrapper