QuLab 2.10.10__cp313-cp313-macosx_10_13_universal2.whl

qulab/sys/device/basedevice.py

@@ -0,0 +1,255 @@
+import copy
+import itertools
+import logging
+import re
+import string
+from collections import defaultdict
+from functools import partial
+from typing import Any, Callable, Literal, NamedTuple
+
+log = logging.getLogger(__name__)
+
+Decorator = Callable[[Callable], Callable]
+
+_buildin_set = set
+
+
+class Source():
+    pass
+
+
+class Sink():
+    pass
+
+
+def action(key: str,
+           method: Literal['get', 'set', 'post', 'delete'] = 'get',
+           **kwds) -> Decorator:
+
+    if any(c in key for c in ",()[]<>"):
+        raise ValueError('Invalid key: ' + key)
+
+    def decorator(func):
+        func.__action__ = key, method, kwds
+        return func
+
+    return decorator
+
+
+def get(key: str, **kwds) -> Decorator:
+    return action(key, 'get', **kwds)
+
+
+def set(key: str, **kwds) -> Decorator:
+    return action(key, 'set', **kwds)
+
+
+def post(key: str, **kwds) -> Decorator:
+    return action(key, 'post', **kwds)
+
+
+def delete(key: str, **kwds) -> Decorator:
+    return action(key, 'delete', **kwds)
+
+
+class _Exclusion(NamedTuple):
+    keys: list
+
+
+def exclude(sections: list):
+    return _Exclusion(sections)
+
+
+def _add_action(attrs: dict, key: str, method: str, func: Callable, doc: dict,
+                sections: dict) -> None:
+    try:
+        mapping = attrs[f'__{method}_actions__']
+    except KeyError:
+        raise ValueError('Invalid method: ' + method)
+    arguments = re.findall(r'\{(\w+)\}', key)
+    doc[method][key] = func.__doc__
+    matrix = {}
+    for arg in arguments:
+        if (arg not in attrs and arg not in sections or arg in sections
+                and arg not in attrs and isinstance(sections[arg], _Exclusion)
+                or arg in sections and arg in attrs
+                and isinstance(sections[arg], _Exclusion) and len(
+                    _buildin_set(attrs[arg]) -
+                    _buildin_set(sections[arg].keys)) == 0):
+            raise ValueError(
+                f'Undefined section: {arg!r} in @action({key!r}, {method!r})')
+        if arg in sections and not isinstance(sections[arg], _Exclusion):
+            matrix[arg] = sections[arg]
+        else:
+            if arg in sections:
+                matrix[arg] = [
+                    k for k in attrs[arg] if k not in sections[arg].keys
+                ]
+            else:
+                matrix[arg] = attrs[arg]
+    for values in itertools.product(*[matrix[arg] for arg in arguments]):
+        kwds = dict(zip(arguments, values))
+        # mapping[key.format(**kwds)] = partial(func, **kwds)
+        mapping[string.Template(key).substitute(kwds)] = partial(func, **kwds)
+
+
+def _build_docs(mapping: dict, attrs: dict) -> str:
+    docs = []
+    for key, doc in mapping.items():
+        if not doc:
+            doc = "No documentation."
+        docs.append(f"key = \"{key}\"")
+        lines = doc.strip().split('\n')
+        docs.extend(lines)
+        docs.append("")
+    return '\n'.join(docs)
+
+
+class DeviceMeta(type):
+
+    def __new__(cls, name, bases, attrs):
+        doc = defaultdict(dict)
+        for method in ['get', 'set', 'post', 'delete']:
+            attrs.setdefault(f'__{method}_actions__', {})
+        for attr in attrs.values():
+            if hasattr(attr, '__action__'):
+                key, method, kwds = attr.__action__
+                _add_action(attrs, key, method, attr, doc, kwds)
+        new_class = super().__new__(cls, name, bases, attrs)
+
+        for method in ['get', 'set', 'post', 'delete']:
+            getattr(new_class,
+                    method).__doc__ = f"{method.upper()}\n\n" + _build_docs(
+                        doc[method], attrs)
+        return new_class
+
+
+class BaseDevice(metaclass=DeviceMeta):
+    __log__ = None
+
+    def __init__(self, address: str = None, **options):
+        self._state = {}
+        self._status = "NOT CONNECTED"
+        self.address = address
+        self.options = options
+
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            pass
+
+    @property
+    def log(self):
+        if self.__log__ is None:
+            self.__log__ = logging.getLogger(
+                f"{self.__class__.__module__}.{self.__class__.__name__}")
+        return self.__log__
+
+    @property
+    def state(self):
+        return copy.deepcopy(self._state)
+    
+    def get_state(self, key: str, default=None) -> Any:
+        return self._state.get(key, default)
+
+    @property
+    def status(self):
+        return self._status
+
+    def open(self) -> None:
+        self._status = "CONNECTED"
+
+    def close(self) -> None:
+        self._status = "NOT CONNECTED"
+        self._state.clear()
+
+    def reset(self) -> None:
+        self._state.clear()
+
+    def get(self, key: str, default: Any = None) -> Any:
+        self.log.info(f'Get {key!r}')
+        if key in self.__get_actions__:
+            result = self.__get_actions__[key](self)
+            self._state[key] = result
+            return result
+        else:
+            return self._state.get(key, default)
+
+    def set(self, key: str, value: Any = None) -> None:
+        self.log.info(f'Set {key!r} = {value!r}')
+        self.__set_actions__[key](self, value)
+        self._state[key] = value
+
+    def post(self, key: str, value: Any = None) -> Any:
+        self.log.info(f'Post {key!r} = {value!r}')
+        return self.__post_actions__[key](self, value)
+
+    def delete(self, key: str) -> None:
+        self.log.info(f'Delete {key!r}')
+        self.__delete_actions__[key](self)
+        del self._state[key]
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.address!r})'
+
+
+class VisaDevice(BaseDevice):
+
+    error_query = 'SYST:ERR?'
+
+    def open(self) -> None:
+        import pyvisa
+        kwds = self.options.copy()
+        if 'backend' in kwds:
+            rm = pyvisa.ResourceManager(kwds.pop('backend'))
+        else:
+            rm = pyvisa.ResourceManager()
+        self.resource = rm.open_resource(self.address, **kwds)
+        self.errors = []
+        super().open()
+
+    def close(self) -> None:
+        super().close()
+        self.resource.close()
+
+    def reset(self) -> None:
+        super().reset()
+        self.resource.write('*RST')
+
+    def check_error(self):
+        if self.error_query:
+            while True:
+                error = self.resource.query(self.error_query)
+                error_code = int(error.split(',')[0])
+                if error_code == 0:
+                    break
+                self.errors.append(error)
+
+    @get('idn')
+    def get_idn(self) -> str:
+        """Get instrument identification."""
+        return self.resource.query('*IDN?')
+
+    @get('opc')
+    def get_opc(self) -> bool:
+        """Get operation complete."""
+        return bool(int(self.resource.query('*OPC?')))
+
+    @get('errors')
+    def get_errors(self) -> list[str]:
+        """Get error queue."""
+        self.check_error()
+        errors = self.errors
+        self.errors = []
+        return errors
+
+    @set('timeout')
+    def set_timeout(self, value: float) -> None:
+        """Set timeout in seconds."""
+        self.resource.timeout = round(value * 1000)
+
+    @get('timeout')
+    def get_timeout(self) -> float:
+        """Get timeout in seconds."""
+        return self.resource.timeout / 1000

qulab/sys/device/loader.py

@@ -0,0 +1,86 @@
+import importlib
+from pathlib import Path
+from typing import Type
+
+from .basedevice import BaseDevice
+
+path = {}
+
+
+def _make_device(module, args, kwds) -> BaseDevice:
+    importlib.reload(module)
+    try:
+        return module.device(*args, **kwds)
+    except AttributeError:
+        return module.Device(*args, **kwds)
+
+
+def create_device_from_file(filepath: str | Path, package_name: str, args,
+                            kwds) -> BaseDevice:
+    """
+    Create a device from a file.
+
+    Parameters
+    ----------
+    filepath : str | Path
+        The path to the driver file.
+    package_name : str
+        The name of the package in which .filepath is located.
+    args : tuple
+        The positional arguments to pass to the device maker.
+    kwds : dict
+        The keyword arguments to pass to the device maker.
+    """
+    filepath = Path(filepath)
+    if not filepath.exists():
+        raise RuntimeError(f"File {filepath} does not exist")
+    module_name = f"{package_name}.{filepath.stem}"
+    spec = importlib.util.spec_from_file_location(module_name, filepath)
+    module = importlib.util.module_from_spec(spec)
+    return _make_device(module, args, kwds)
+
+
+def create_device_from_module(module_name: str, args, kwds) -> BaseDevice:
+    """
+    Create a device from a module.
+
+    Parameters
+    ----------
+    module_name : str
+        The name of the module to import.
+    args : tuple
+        The positional arguments to pass to the device maker.
+    kwds : dict
+        The keyword arguments to pass to the device maker.
+    """
+    module = importlib.import_module(module_name)
+    return _make_device(module, args, kwds)
+
+
+def create_device(driver_name: str, *args, **kwds) -> BaseDevice:
+    """
+    Create a device from a driver.
+
+    Parameters
+    ----------
+    driver_name : str
+        The name of the driver to use.
+    args : tuple
+        The positional arguments to pass to the device maker.
+    kwds : dict
+        The keyword arguments to pass to the device maker.
+    """
+    try:
+        for package_name, p in path.items():
+            try:
+                dev = create_device_from_file(
+                    Path(p) / f"{driver_name}.py", package_name, args, kwds)
+                return dev
+            except:
+                pass
+
+        dev = create_device_from_module(f"qulab.sys.drivers.{driver_name}",
+                                        args, kwds)
+        return dev
+    except:
+        raise RuntimeError(f"Can not find driver {driver_name!r}")

qulab/sys/device/utils.py

@@ -0,0 +1,79 @@
+import struct
+
+import numpy as np
+
+types = {
+    "b"      : (  int, 'b'), "B"      : (  int, 'B'),
+    "h"      : (  int, 'h'), "H"      : (  int, 'H'),
+    "i"      : (  int, 'i'), "I"      : (  int, 'I'),
+    "q"      : (  int, 'q'), "Q"      : (  int, 'Q'),
+    "f"      : (float, 'f'), "d"      : (float, 'd'),
+    "int8"   : (  int, 'b'), "uint8"  : (  int, 'B'),
+    "int16"  : (  int, 'h'), "uint16" : (  int, 'H'),
+    "int32"  : (  int, 'i'), "uint32" : (  int, 'I'),
+    "int64"  : (  int, 'q'), "uint64" : (  int, 'Q'),
+    "float"  : (float, 'f'), "double" : (float, 'd'),
+    "float32": (float, 'f'), "float64": (float, 'd')
+} # yapf: disable
+
+
+def IEEE_488_2_BinBlock(datalist, dtype="int16", is_big_endian=True):
+    """
+    将一组数据打包成 IEEE 488.2 标准二进制块
+
+    Args:
+        datalist : 要打包的数字列表
+        dtype    : 数据类型
+        endian   : 字节序
+
+    Returns:
+        binblock, header
+        二进制块, 以及其 'header'
+    """
+    if isinstance(datalist, bytes):
+        datablock = datalist
+    else:
+        datalist = np.asarray(datalist)
+        datalist.astype(types[dtype][0])
+        if is_big_endian:
+            endianc = '>'
+        else:
+            endianc = '<'
+        datablock = struct.pack(
+            '%s%d%s' % (endianc, len(datalist), types[dtype][1]), *datalist)
+    size = '%d' % len(datablock)
+    header = '#%d%s' % (len(size), size)
+
+    return header.encode() + datablock
+
+
+def decode_IEEE_488_2_BinBlock(data, dtype="int16", is_big_endian=True):
+    """
+    解析 IEEE 488.2 标准二进制块
+
+    Args:
+        data : 二进制块
+        dtype : 数据类型
+        endian : 字节序
+
+    Returns:
+        datalist
+    """
+    if isinstance(data, bytes):
+        header = data[:2]
+        size = int(data[2:2 + int(header[1:2])])
+        assert size == len(data) - 2 - int(header[1:2]), "data size error"
+        datablock = data[2 + int(header[1:2]):size + 2 + int(header[1:2])]
+    else:
+        raise ValueError("data must be bytes")
+
+    if is_big_endian:
+        endianc = '>'
+    else:
+        endianc = '<'
+    datalist = list(
+        struct.unpack(
+            '%s%d%s' % (endianc, size // struct.calcsize(types[dtype][1]),
+                        types[dtype][1]), datablock))
+
+    return datalist

qulab/sys/drivers/FakeInstrument.py

@@ -0,0 +1,68 @@
+import numpy as np
+
+from qulab.sys import BaseDevice, exclude, get, set
+
+
+class Device(BaseDevice):
+
+    channel = [
+        'X1', 'X2', 'X3', 'X4', 'Y1', 'Y2', 'Y3', 'Y4', 'Z1', 'Z2', 'Z3', 'Z4',
+        'M1', 'M2'
+    ]
+
+    @get('IDN')
+    def get_idn(self) -> str:
+        return 'Fake Instrument'
+
+    @get('${channel_type}${channel_num}.SampleRate',
+         channel_type=['X', 'Y'],
+         channel_num=[1, 2, 3, 4])
+    def get_sample_rate_DDS(self, channel_type: str,
+                            channel_num: int) -> float:
+        return 6e9
+
+    @get('${channel}.SampleRate', channel=['Z1', 'Z2', 'Z3', 'Z4'])
+    def get_sample_rate_Z(self, channel: str) -> float:
+        return 2e9
+
+    @get('${channel}.SampleRate', channel=['M1', 'M2'])
+    def get_sample_rate_AD(self, channel: str) -> float:
+        return 1e9
+
+    @set('${channel}.Vpp', channel=exclude(['M1', 'M2']))
+    def set_voltage(self, value: float, channel: str) -> None:
+        self.log.info(f'Set {channel} Vpp to {value}')
+
+    # @get('${channel}.Vpp', channel=exclude(['M1', 'M2']))
+    # def get_voltage(self, channel: str, default=0.0) -> float:
+    #     return self._state.get(f'{channel}.Vpp', default)
+
+    @set('${channel}.Offset', channel=exclude(['M1', 'M2']))
+    def set_offset(
+        self,
+        value: float,
+        channel: str,
+    ) -> None:
+        self.log.info(f'Set {channel} offset to {value}')
+
+    # @get('${channel}.Offset', channel=exclude(['M1', 'M2']))
+    # def get_offset(self, channel: str, default=0.0) -> float:
+    #     return self._state.get(f'{channel}.Offset', default)
+
+    @set('${channel}.Waveform', channel=exclude(['M1', 'M2']))
+    def set_waveform(self, value, channel: str) -> None:
+        self.log.info(f'Set {channel} waveform to {value!r}')
+
+    # @get('${channel}.Waveform', channel=exclude(['M1', 'M2']))
+    # def get_waveform(self, channel: str, default=None) -> str:
+    #     return self._state.get(f'{channel}.Waveform', default)
+
+    @set('${channel}.Size', channel=['M1', 'M2'])
+    def set_size(self, value, channel: str) -> None:
+        self.log.info(f'Set {channel} size to {value!r}')
+
+    @get('${channel}.Trace', channel=['M1', 'M2'])
+    def get_random_data(self, channel: str) -> np.ndarray:
+        size = self._state.get(f'{channel}.Size', 1024)
+        shots = self._state.get(f'{channel}.Shots', 128)
+        return np.random.randn(shots, size)

qulab/sys/ipy_events.py

@@ -0,0 +1,125 @@
+import contextlib
+import time
+from datetime import datetime
+
+from .storage.crud import create_cell, create_notebook
+from .storage.models import Notebook
+
+__current_notebook_id = None
+
+
+def _get_current_notebook(session):
+    global __current_notebook_id
+    if __current_notebook_id is None:
+        notebook = create_notebook(session, 'Untitled')
+        session.commit()
+        __current_notebook_id = notebook.id
+    else:
+        notebook = session.query(Notebook).filter(
+            Notebook.id == __current_notebook_id).one()
+    return notebook
+
+
+def _get_current_cell(session):
+    notebook = _get_current_notebook(session)
+    if notebook.cells:
+        return notebook.cells[-1]
+    else:
+        return None
+
+
+def _save_inputCells(session, inputCells=None):
+    notebook = _get_current_notebook(session)
+    aready_saved = len(notebook.cells)
+    if inputCells:
+        for cell in inputCells[aready_saved:]:
+            create_cell(session, notebook, cell)
+
+
+def _update_inputCells(session):
+    cell = _get_current_cell(session)
+    if cell:
+        cell.ftime = datetime.utcnow()
+        session.add(cell)
+
+
+__sessionmaker = None
+
+
+@contextlib.contextmanager
+def get_session():
+    global __session
+    if __sessionmaker is None:
+        raise RuntimeError('Sessionmaker is not set')
+    session = __sessionmaker()
+    try:
+        yield session
+        session.commit()
+    except:
+        session.rollback()
+
+
+def set_sessionmaker(sessionmaker):
+    global __sessionmaker
+    __sessionmaker = sessionmaker
+
+
+def get_current_notebook(session=None):
+    try:
+        if session is None:
+            with get_session() as session:
+                return _get_current_notebook(session)
+        else:
+            return _get_current_notebook(session)
+    except RuntimeError:
+        return None
+
+
+def get_current_cell_id(session=None):
+    try:
+        if session is None:
+            with get_session() as session:
+                return _get_current_cell(session).id
+        else:
+            return _get_current_cell(session).id
+    except RuntimeError as e:
+        return None
+
+
+def get_inputCells():
+    try:
+        from IPython import get_ipython
+
+        return get_ipython().user_ns['In']
+    except Exception as e:
+        # raise e
+        return ['']
+
+
+def autosave_cells(ipython):
+    try:
+        with get_session() as session:
+            _save_inputCells(session, get_inputCells())
+    except:
+        pass
+
+
+def update_timestamp(ipython):
+    try:
+        with get_session() as session:
+            _update_inputCells(session)
+    except:
+        pass
+
+
+def setup_ipy_events():
+    try:
+        from IPython import get_ipython
+
+        ipython = get_ipython()
+
+        if ipython is not None:
+            ipython.events.register('pre_run_cell', autosave_cells)
+            ipython.events.register('post_run_cell', update_timestamp)
+    except:
+        pass