177d8777c1
All checks were successful
Run apocalypse tests / Explore-Gitea-Actions (push) Successful in 10s
122 lines
2.9 KiB
Python
122 lines
2.9 KiB
Python
#copy https://github.com/paulscherrerinstitute/sf_datafiles/blob/0598589492ef4214f867df39ad5e87480bbb4b94/sfdata/load.py
|
|
import collections
|
|
import itertools
|
|
import os.path
|
|
from glob import iglob
|
|
|
|
|
|
def make_loader(instrument="*a", pgroup="p*", run="*"):
|
|
|
|
def loader_wrapper(instrument=instrument, pgroup=pgroup, run=run, debug=False):
|
|
runs = load_many(instrument, pgroup, run, debug)
|
|
if debug:
|
|
yield from runs
|
|
else:
|
|
for run in runs:
|
|
fn = os.path.join(run, "meta/scan.json")
|
|
|
|
return loader_wrapper
|
|
|
|
|
|
def load_many(instrument, pgroup, run, debug):
|
|
for i, p, r in product(instrument, pgroup, run):
|
|
yield from load_one(i, p, r, debug)
|
|
|
|
|
|
def product(*vals):
|
|
vals = ensure_seqs(*vals)
|
|
return itertools.product(*vals)
|
|
|
|
|
|
def ensure_seqs(*vals):
|
|
return (ensure_seq(i) for i in vals)
|
|
|
|
|
|
def ensure_seq(val):
|
|
if isiterable(val):
|
|
return val
|
|
return (val, )
|
|
|
|
|
|
def isiterable(val):
|
|
return isinstance(val, collections.abc.Iterable) and not isinstance(val, str)
|
|
|
|
|
|
def load_one(instrument, pgroup, run, debug):
|
|
pgroup = harmonize(pgroup, "p", 5)
|
|
run = harmonize(run, "run", 4)
|
|
pattern = f"/sf/{instrument}/data/{pgroup}/raw/{run}"
|
|
res = sorted(iglob(pattern))
|
|
if debug:
|
|
print(pattern, "-> #runs:", len(res))
|
|
return [pattern]
|
|
else:
|
|
return res
|
|
|
|
|
|
def harmonize(val, prefix, nints):
|
|
val = str(val)
|
|
val = strip_prefix(val, prefix)
|
|
val = pad_if_int(val, nints)
|
|
val = prefix + val
|
|
return val
|
|
|
|
|
|
def strip_prefix(string, prefix):
|
|
if string.startswith(prefix):
|
|
n = len(prefix)
|
|
string = string[n:]
|
|
return string
|
|
|
|
|
|
def pad_if_int(val, length):
|
|
try:
|
|
val = int(val)
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
val = str(val)
|
|
val = val.zfill(4)
|
|
return val
|
|
|
|
|
|
# create a default load
|
|
load = make_loader()
|
|
|
|
if __name__ == "__main__":
|
|
from functools import partial
|
|
|
|
iload = partial(load, debug=True)
|
|
load = lambda *args, **kwargs: list(iload(*args, **kwargs))
|
|
|
|
assert load(instrument="TEST") == ["/sf/TEST/data/p*/raw/run*"]
|
|
assert load(pgroup="TEST") == ["/sf/*a/data/pTEST/raw/run*"]
|
|
assert load(run="TEST") == ["/sf/*a/data/p*/raw/runTEST"]
|
|
|
|
ref = ["/sf/*a/data/p*/raw/run*"]
|
|
|
|
assert ref == load()
|
|
assert ref == load(run="*")
|
|
assert ref == load(run="run*")
|
|
assert ref == load(pgroup="*")
|
|
assert ref == load(pgroup="p*")
|
|
|
|
ref = ["/sf/*a/data/p*/raw/run0001"]
|
|
|
|
assert ref == load(run=1)
|
|
assert ref == load(run="1")
|
|
assert ref == load(run="01")
|
|
assert ref == load(run="run1")
|
|
assert ref == load(run="run00001")
|
|
|
|
ref = ["/sf/*a/data/p12345/raw/run*"]
|
|
|
|
assert ref == load(pgroup=12345)
|
|
assert ref == load(pgroup="12345")
|
|
assert ref == load(pgroup="p12345")
|
|
|
|
ref = ["/sf/*a/data/p*/raw/run0010", "/sf/*a/data/p*/raw/run0020"]
|
|
|
|
assert ref == load(run=[10, 20])
|
|
assert ref == load(run=iter(range(10, 20 + 10, 10)))
|