diff --git a/ptgctl/tools/display.py b/ptgctl/tools/display.py index 2bfa329..9355a6a 100644 --- a/ptgctl/tools/display.py +++ b/ptgctl/tools/display.py @@ -1,9 +1,12 @@ '''This contains functions to display values from the API in interesting ways. ''' +import os +import contextlib import io import json as json_ import time +import tqdm import asyncio import datetime import numpy as np @@ -115,6 +118,24 @@ async def raw(api, stream_id, utf=False, **kw): print("could not decode:", data) +@util.async2sync +@util.interruptable +async def file(api, stream_id, out_dir, include_timestamps=False, **kw): + async with api.data_pull_connect(stream_id, **kw) as ws: + with contextlib.ExitStack() as stack: + files = {} + pbars = {} + while True: + for sid, ts, data in await ws.recv_data(): + if sid not in files: + files[sid] = stack.enter_context(open(os.path.join(out_dir, f'{sid}.txt'), 'w')) + pbars[sid] = tqdm.tqdm(desc=sid) + if include_timestamps: + files[sid].write(f'{ts}:') + files[sid].write(f"{data.decode('utf-8')}\n") + pbars[sid].update() + pbars[sid].set_description(f'{sid}: {str(data)[:20]}') + @util.async2sync @util.interruptable diff --git a/ptgctl/util/cli.py b/ptgctl/util/cli.py index 50c31eb..35eeef1 100644 --- a/ptgctl/util/cli.py +++ b/ptgctl/util/cli.py @@ -4,6 +4,28 @@ +class CachedProperty: + def __init__(self, getattr): + super().__init__() + self.new = getattr + self.key = '__{}'.format(getattr.__name__) + self._blank = None + + # Works as a property that instantiates a nested class on first access. + # the created instance will be used for subsequent access. + def __get__(self, instance, owner=None): + if instance is None: + if self._blank is None: + self._blank = self.new(self) + return self._blank + try: + return getattr(instance, self.key) + except AttributeError: + x = self.new(instance) + setattr(instance, self.key, x) + return x + + class BoundModule: '''Acts as a method namespace for a class using a module as the namespace. @@ -96,17 +118,10 @@ def greetings(self): import greetings return greetings ''' - attr = f'__{get_module.__name__}' - @property + @CachedProperty @functools.wraps(get_module) def inner(self): - # only create the bound module once - try: - return getattr(self, attr) - except AttributeError: - bm = cls(self, get_module) - setattr(self, attr, bm) - return bm + return cls(self, get_module) return inner bound_module = BoundModule._bind