QuLab 2.0.1__cp311-cp311-macosx_10_9_universal2.whl → 2.0.3__cp311-cp311-macosx_10_9_universal2.whl

qulab/scan/utils.py

@@ -1,37 +1,83 @@
+import ast
+import asyncio
 import inspect
-from concurrent.futures import Future
-
-
-def call_func_with_kwds(func, args, kwds, log=None):
-    funcname = getattr(func, '__name__', repr(func))
-    sig = inspect.signature(func)
-    for p in sig.parameters.values():
-        if p.kind == p.VAR_KEYWORD:
-            return func(*args, **kwds)
-    kw = {
-        k: v
-        for k, v in kwds.items()
-        if k in list(sig.parameters.keys())[len(args):]
-    }
+from typing import Any, Callable
+
+from .expression import Env, Expression
+
+
+def is_valid_identifier(s: str) -> bool:
+    """
+    Check if a string is a valid identifier.
+    """
+    try:
+        ast.parse(f"f({s}=0)")
+        return True
+    except SyntaxError:
+        return False
+
+
+async def async_next(aiter):
+    try:
+        if hasattr(aiter, '__anext__'):
+            return await aiter.__anext__()
+        else:
+            return next(aiter)
+    except StopIteration:
+        raise StopAsyncIteration from None
+
+
+async def async_zip(*aiters):
+    aiters = [
+        ait.__aiter__() if hasattr(ait, '__aiter__') else iter(ait)
+        for ait in aiters
+    ]
+    try:
+        while True:
+            # 使用 asyncio.gather 等待所有异步生成器返回下一个元素
+            result = await asyncio.gather(*(async_next(ait) for ait in aiters))
+            yield tuple(result)
+    except StopAsyncIteration:
+        # 当任一异步生成器耗尽时停止迭代
+        return
+
+
+async def call_function(func: Callable | Expression, variables: dict[str,
+                                                                     Any]):
+    if isinstance(func, Expression):
+        env = Env()
+        for name in func.symbols():
+            if name in variables:
+                if inspect.isawaitable(variables[name]):
+                    variables[name] = await variables[name]
+                env.variables[name] = variables[name]
+            else:
+                raise ValueError(f'{name} is not provided.')
+        return func.eval(env)
+
     try:
-        args = [
-            arg.result() if isinstance(arg, Future) else arg for arg in args
-        ]
-        kw = {
-            k: v.result() if isinstance(v, Future) else v
-            for k, v in kw.items()
-        }
-        return func(*args, **kw)
+        sig = inspect.signature(func)
     except:
-        if log:
-            log.exception(f'Call {funcname} with {args} and {kw}')
-        raise
-    finally:
-        if log:
-            log.debug(f'Call {funcname} with {args} and {kw}')
-
-
-def try_to_call(x, args, kwds, log=None):
-    if callable(x):
-        return call_func_with_kwds(x, args, kwds, log)
-    return x
+        return func()
+    args = []
+    for name, param in sig.parameters.items():
+        if param.kind == param.POSITIONAL_OR_KEYWORD:
+            if name in variables:
+                if inspect.isawaitable(variables[name]):
+                    variables[name] = await variables[name]
+                args.append(variables[name])
+            elif param.default is not param.empty:
+                args.append(param.default)
+            else:
+                raise ValueError(f'parameter {name} is not provided.')
+        elif param.kind == param.VAR_POSITIONAL:
+            raise ValueError('not support VAR_POSITIONAL')
+        elif param.kind == param.VAR_KEYWORD:
+            ret = func(**variables)
+            if inspect.isawaitable(ret):
+                ret = await ret
+            return ret
+    ret = func(*args)
+    if inspect.isawaitable(ret):
+        ret = await ret
+    return ret

qulab/sys/rpc/zmq_socket.py

@@ -0,0 +1,209 @@
+from typing import Optional
+
+import zmq
+import zmq.asyncio
+import zmq.auth
+from watchdog.events import FileSystemEventHandler
+from watchdog.observers import Observer
+from zmq.auth.asyncio import AsyncioAuthenticator
+from zmq.auth.thread import ThreadAuthenticator
+
+
+class ReloadCertificatesHandler(FileSystemEventHandler):
+
+    def __init__(self, zmq: 'ZMQContextManager'):
+        self.zmq = zmq
+
+    def on_modified(self, event):
+        self.zmq.reload_certificates()
+
+
+class ZMQContextManager:
+    """
+    A context manager for managing ZeroMQ sockets with asynchronous support.
+    It handles the creation, connection, binding, and security configuration of
+    the sockets and ensures proper resource cleanup.
+
+    The context manager can be used as a synchronous context manager or an
+    asynchronous context manager. When used as an asynchronous context manager,
+    the socket is created with an asyncio context and can be used with the
+    `await` keyword.
+
+    The security settings for the socket can be configured using the secret key
+    and public key parameters. If the secret key is provided, the socket is
+    configured to use ZeroMQ curve encryption. The public key of the server can
+    also be provided for client sockets to connect to a server with curve
+    encryption. You can also provide the paths to the secret key and public key
+    files, and the public key of the server to load the keys from files. The
+    keys can be reloaded automatically when the files are modified by setting
+    the `public_keys_location` parameter.
+
+    To generate a secret key and public key pair, you can use the following
+    commands:
+    
+        ```bash
+        # Generate a secret key and public key pair
+        python -c "import zmq.auth; zmq.auth.create_certificates('.', 'filename')"
+        ```
+
+    Attributes:
+        socket_type: zmq.SocketType
+            The type of the socket to create (e.g., zmq.REP).
+        bind: str, optional
+            The address to bind the socket to.
+        connect: str, optional
+            The address to connect the socket to.
+        secret_key_file: str, optional
+            The path to the secret key file for ZeroMQ curve encryption.
+        server_public_key_file: str, optional
+            The path to the public key file for the server in ZeroMQ curve
+            encryption.
+        public_keys_location: str, optional
+            The location to store the public keys for ZeroMQ curve encryption.
+        secret_key: bytes, optional
+            The secret key for ZeroMQ curve encryption.
+        public_key: bytes, optional
+            The public key for ZeroMQ curve encryption.
+        server_public_key: bytes, optional
+            The public key for the server in ZeroMQ curve encryption.
+
+    Methods:
+        _create_socket: zmq.Socket
+            Creates and configures a ZeroMQ socket.
+        _close_socket: None
+            Closes the ZeroMQ socket and the context, and stops the authenticator
+            if it was started.
+
+    Examples:
+        Create a REP socket and bind it to an address:
+
+        >>> async with ZMQContextManager(zmq.REP, bind='tcp://*:5555') as socket:
+        ...     while True:
+        ...         message = await socket.recv()
+        ...         await socket.send(message)
+
+        Create a REQ socket and connect it to an address:
+
+        >>> with ZMQContextManager(zmq.REQ, connect='tcp://localhost:5555') as socket:
+        ...     socket.send(b'Hello')
+        ...     message = socket.recv()
+    """
+
+    def __init__(self,
+                 socket_type: zmq.SocketType,
+                 bind: Optional[str] = None,
+                 connect: Optional[str] = None,
+                 secret_key_file: Optional[str] = None,
+                 server_public_key_file: Optional[str] = None,
+                 public_keys_location: Optional[str] = None,
+                 secret_key: Optional[bytes] = None,
+                 public_key: Optional[bytes] = None,
+                 server_public_key: Optional[bytes] = None):
+        self.socket_type = socket_type
+        if bind is None and connect is None:
+            raise ValueError("Either 'bind' or 'connect' must be specified.")
+        if bind is not None and connect is not None:
+            raise ValueError("Both 'bind' and 'connect' cannot be specified.")
+        self.bind = bind
+        self.connect = connect
+        self.secret_key = secret_key
+        self.public_key = public_key
+        self.server_public_key = server_public_key
+
+        if secret_key_file:
+            self.public_key, self.secret_key = zmq.auth.load_certificate(
+                secret_key_file)
+
+        if (self.secret_key is not None and self.public_key is None
+                or self.secret_key is None and self.public_key is not None):
+            raise ValueError(
+                "Both secret key and public key must be specified.")
+
+        if server_public_key_file:
+            self.server_public_key = zmq.auth.load_certificate(
+                server_public_key_file)[0]
+
+        self.public_keys_location = public_keys_location
+
+        self.observer = None
+        self.auth = None
+        self.context = None
+        self.socket = None
+
+    def _create_socket(self, asyncio=False) -> zmq.Socket:
+        """
+        Creates and configures a ZeroMQ socket. Sets up security if required,
+        and binds or connects the socket according to the specified settings.
+
+        Returns:
+            zmq.Socket: The configured ZeroMQ socket.
+        """
+        if asyncio:
+            self.context = zmq.asyncio.Context()
+        else:
+            self.context = zmq.Context.instance()
+
+        self.socket = self.context.socket(self.socket_type)
+        self.auth = None
+
+        if self.bind and self.secret_key:
+            if asyncio:
+                self.auth = AsyncioAuthenticator(self.context)
+            else:
+                self.auth = ThreadAuthenticator(self.context)
+            self.auth.start()
+            self.reload_certificates()
+            self.auto_reload_certificates()
+            self.socket.curve_server = True  # must come before bind
+
+        if self.secret_key:
+            self.socket.curve_secretkey = self.secret_key
+            self.socket.curve_publickey = self.public_key
+
+        if self.bind:
+            self.socket.bind(self.bind)
+        if self.connect:
+            if self.server_public_key:
+                self.socket.curve_serverkey = self.server_public_key
+            self.socket.connect(self.connect)
+        return self.socket
+
+    def reload_certificates(self):
+        if self.public_keys_location and self.auth:
+            self.auth.configure_curve(domain='*',
+                                      location=self.public_keys_location)
+
+    def auto_reload_certificates(self):
+        self.observer = Observer()
+        self.observer.schedule(ReloadCertificatesHandler(self),
+                               self.public_keys_location,
+                               recursive=False)
+        self.observer.start()
+
+    def _close_socket(self) -> None:
+        """
+        Closes the ZeroMQ socket and the context, and stops the authenticator
+        if it was started.
+        """
+        if self.observer:
+            self.observer.stop()
+            self.observer.join()
+        self.socket.close()
+        if self.auth:
+            self.auth.stop()
+        self.context.term()
+
+    def __enter__(self) -> zmq.Socket:
+        return self._create_socket(asyncio=False)
+
+    def __exit__(self, exc_type: Optional[type], exc_val: Optional[Exception],
+                 exc_tb: Optional[type]) -> None:
+        self._close_socket()
+
+    async def __aenter__(self) -> zmq.Socket:
+        return self._create_socket(asyncio=True)
+
+    async def __aexit__(self, exc_type: Optional[type],
+                        exc_val: Optional[Exception],
+                        exc_tb: Optional[type]) -> None:
+        self._close_socket()

qulab/version.py

@@ -1 +1 @@
-__version__ = "2.0.1"
+__version__ = "2.0.3"

qulab/visualization/_autoplot.py

@@ -2,7 +2,7 @@ import math
 
 import matplotlib.pyplot as plt
 import numpy as np
-from matplotlib.colors import LogNorm
+from matplotlib.colors import LogNorm, Normalize, SymLogNorm
 from matplotlib.ticker import EngFormatter, LogFormatterSciNotation
 from scipy.interpolate import griddata
 
@@ -239,10 +239,16 @@ def plot_img(x,
     kwds.setdefault('aspect', 'auto')
     kwds.setdefault('interpolation', 'nearest')
 
+    vmin = kwds.get('vmin', np.nanmin(z))
+    vmax = kwds.get('vmax', np.nanmax(z))
     if zscale == 'log':
-        vmim = kwds.get('vmin', np.min(z))
-        vmax = kwds.get('vmax', np.max(z))
-        kwds.setdefault('norm', LogNorm(vmax=vmax, vmin=vmim))
+        kwds.setdefault('norm', LogNorm(vmax=vmax, vmin=vmin))
+    elif zscale == 'symlog':
+        kwds.setdefault('norm', SymLogNorm(vmax=vmax,
+                                           vmin=vmin,
+                                           linthresh=1e-5))
+    else:
+        kwds.setdefault('norm', Normalize(vmin=vmin, vmax=vmax))
     zlabel = f"{zlabel} [{z_unit}]" if z_unit else zlabel
 
     band_area = False
@@ -303,7 +309,7 @@ def plot_img(x,
         pass
     ax.set_xlabel(xlabel)
     ax.set_ylabel(ylabel)
-    cb = fig.colorbar(img, ax=ax)
+    cb = fig.colorbar(img, ax=ax, norm=kwds.get('norm', None))
     cb.set_label(zlabel)