diff --git a/CHANGELOG.md b/CHANGELOG.md index f43f7e3..2c5fd54 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ ## [1.15.2] - 2019-06-07 ### Added - Store unique stream ID inside the `["info"]["stream_id"]` dict value ([#19](https://github.com/xdf-modules/xdf-Python/pull/19) by [Clemens Brunner](https://github.com/cbrnr)). +- Add option to load only specific streams ([#24](https://github.com/xdf-modules/xdf-Python/pull/24) by [Clemens Brunner](https://github.com/cbrnr)). ## [1.15.1] - 2019-04-26 ### Added diff --git a/pyxdf/__init__.py b/pyxdf/__init__.py index ad41151..51968a7 100644 --- a/pyxdf/__init__.py +++ b/pyxdf/__init__.py @@ -1,12 +1,12 @@ # Authors: Christian Kothe & the Intheon pyxdf team +# Clemens Brunner # # License: BSD (2-clause) from pkg_resources import get_distribution, DistributionNotFound try: __version__ = get_distribution(__name__).version -except DistributionNotFound: - # package is not installed +except DistributionNotFound: # package is not installed __version__ = None -from .pyxdf import load_xdf +from .pyxdf import load_xdf, resolve_streams, match_streaminfos diff --git a/pyxdf/pyxdf.py b/pyxdf/pyxdf.py index 4dee842..6a730cb 100644 --- a/pyxdf/pyxdf.py +++ b/pyxdf/pyxdf.py @@ -67,6 +67,7 @@ def __init__(self, xml): def load_xdf(filename, + select_streams=None, on_chunk=None, synchronize_clocks=True, handle_clock_resets=True, @@ -94,6 +95,20 @@ def load_xdf(filename, Args: filename : name of the file to import (*.xdf or *.xdfz) + select_streams : int | list[int] | list[dict] | None + One or more stream IDs to load. Accepted values are: + - int or list[int]: load only specified stream IDs, e.g. + select_streams=5 loads only the stream with stream ID 5, whereas + select_streams=[2, 4] loads only streams with stream IDs 2 and 4. + - list[dict]: load only streams matching the query, e.g. + select_streams=[{'type': 'EEG'}] loads all streams of type 'EEG'. + Entries within a dict must all match a stream, e.g. + select_streams=[{'type': 'EEG', 'name': 'TestAMP'}] matches streams + with both type 'EEG' *and* name 'TestAMP'. If + select_streams=[{'type': 'EEG'}, {'name': 'TestAMP'}], streams + matching either the type *or* the name will be loaded. + - None: load all streams (default). + synchronize_clocks : Whether to enable clock synchronization based on ClockOffset chunks. (default: true) @@ -155,10 +170,10 @@ def load_xdf(filename, streams : list of dicts, one for each stream; the dicts have the following content: ['time_series'] entry: contains the stream's time series - [#Channels x #Samples] this matrix is of the type declared in - ['info']['channel_format'] - ['time_stamps'] entry: contains the time stamps for each sample - (synced across streams) + [#Channels x #Samples] this matrix is of the type declared + in ['info']['channel_format'] + ['time_stamps'] entry: contains the time stamps for each + sample (synced across streams) ['info'] field: contains the meta-data of the stream (all values are strings) @@ -179,14 +194,31 @@ def load_xdf(filename, Examples: load the streams contained in a given XDF file - >>> streams, fileheader = load_xdf('C:\Recordings\myrecording.xdf') + >>> streams, fileheader = load_xdf('myrecording.xdf') """ logger.info('Importing XDF file %s...' % filename) if not os.path.exists(filename): raise Exception('file %s does not exist.' % filename) - # dict of returned streams, in order of apparance, indexed by stream id + # if select_streams is an int or a list of int, load only streams + # associated with the corresponding stream IDs + # if select_streams is a list of dicts, use this to query and load streams + # associated with these properties + if select_streams is None: + pass + elif isinstance(select_streams, int): + select_streams = [select_streams] + elif all([isinstance(elem, dict) for elem in select_streams]): + select_streams = match_streaminfos(resolve_streams(filename), + select_streams) + if not select_streams: # no streams found + raise ValueError("No matching streams found.") + elif not all([isinstance(elem, int) for elem in select_streams]): + raise ValueError("Argument 'select_streams' must be an int, a list of " + "ints or a list of dicts.") + + # dict of returned streams, in order of appearance, indexed by stream id streams = OrderedDict() # dict of per-stream temporary data (StreamData), indexed by stream id temp = {} @@ -195,22 +227,9 @@ def load_xdf(filename, # number of bytes in the file for fault tolerance filesize = os.path.getsize(filename) - # read file contents ([SomeText] below refers to items in the XDF Spec) - filename = Path(filename) # convert to pathlib object - if filename.suffix == '.xdfz' or filename.suffixes == ['.xdf', '.gz']: - f_open = gzip.open - else: - f_open = open - - with f_open(filename, 'rb') as f: - # read [MagicCode] - if f.read(4) != b'XDF:': - raise Exception('not a valid XDF file: %s' % filename) - - # for each chunk... - StreamId = None + with open_xdf(filename) as f: + # for each chunk while True: - # noinspection PyBroadException try: # read [NumLengthBytes], [Length] @@ -231,9 +250,16 @@ def load_xdf(filename, if tag in [2, 3, 4, 6]: StreamId = struct.unpack('