QuLab 2.10.10__cp313-cp313-win_amd64.whl

qulab/scan/query.py

@@ -0,0 +1,387 @@
+import functools
+import math
+from pathlib import Path
+
+import dill
+import ipywidgets as widgets
+import zmq
+from IPython.display import display
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+
+from qulab.sys.rpc.zmq_socket import ZMQContextManager
+
+from .record import Record
+from .scan import default_server
+from .server import get_local_record
+
+
+def get_record(id,
+               database=default_server,
+               socket=None,
+               session=None) -> Record:
+    if isinstance(database, str) and database.startswith('tcp://'):
+        with ZMQContextManager(zmq.DEALER, connect=database,
+                               socket=socket) as socket:
+            socket.send_pyobj({
+                'method': 'record_description',
+                'record_id': id
+            })
+            d = dill.loads(socket.recv_pyobj())
+            if d is None:
+                raise ValueError(f'No record with id {id}')
+            d.id = id
+            d.database = database
+            d._file = None
+            d._sock = socket
+            return d
+    else:
+        if session is None:
+            db_file = Path(database) / 'data.db'
+            engine = create_engine(f'sqlite:///{db_file}')
+            Session = sessionmaker(bind=engine)
+            with Session() as session:
+                return get_local_record(session, id, database)
+        else:
+            return get_local_record(session, id, database)
+
+
+def load_record(file):
+    return Record.load(file)
+
+
+def _format_tag(tag):
+    if tag.startswith('!'):
+        return f'<code style="color: white; background: red">{tag}</code>'
+    elif tag.startswith('?'):
+        return f'<code style="background: orange">{tag}</code>'
+    elif tag.startswith('@'):
+        return f'<code style="color: white; background: green">{tag}</code>'
+    elif tag.startswith('#'):
+        return f'<code style="color: white; background: blue">{tag}</code>'
+    elif tag.startswith('$'):
+        return f'<code style="color: white; background: purple">{tag}</code>'
+    elif tag.startswith('%'):
+        return f'<code style="color: white; background: brown">{tag}</code>'
+    else:
+        return f'<code>{tag}</code>'
+
+
+def _format_value(title, value):
+    if title in ['ID', 'App']:
+        return str(value)
+    elif title in ['created time']:
+        return str(value).split('.')[0]
+    elif title in ['tags']:
+        tags = sorted(value)
+        if len(tags) <= 6:
+            tags = [_format_tag(t) for t in tags]
+        elif len(tags) <= 12:
+            tags = [_format_tag(t) for t in tags[:6]
+                    ] + ['<br />'] + [_format_tag(t) for t in tags[6:]]
+        elif len(tags) <= 18:
+            tags = [_format_tag(t) for t in tags[:6]] + ['<br />'] + [
+                _format_tag(t) for t in tags[6:12]
+            ] + ['<br />'] + [_format_tag(t) for t in tags[12:]]
+        else:
+            tags = [_format_tag(t) for t in tags[:6]] + ['<br />'] + [
+                _format_tag(t) for t in tags[6:12]
+            ] + ['<br />'] + [_format_tag(t) for t in tags[12:17]] + ['...']
+        return ' '.join(tags)
+    else:
+        return repr(value)
+
+
+def _format_table(table):
+    tb = [
+        """<div class="output_subarea output_markdown rendered_html" dir="auto">"""
+        """<table>"""
+        """<thead>"""
+        """<tr>"""
+    ]
+    for s in table['header']:
+        tb.append(f'<th style="text-align: center">{s}</th>')
+    tb.append("""</tr></thead><tbody>""")
+
+    for row in table["body"]:
+        r = ["<tr>"]
+        for k, v in zip(table['header'], row):
+            s = _format_value(k, v)
+            r.append(f'<td style="text-align: center">{s}</td>')
+        r.append("</tr>")
+        tb.append(''.join(r))
+    tb.append("""</tbody></table>"""
+              """</div>""")
+    return ''.join(tb)
+
+
+class _QueryState():
+
+    def __init__(self):
+        self.table = None
+        self.apps = {}
+        self.page = 1
+        self.total = 0
+        self.app_prefix = []
+        self.app_name = '*'
+        self.records_per_page = 10
+
+    def list_apps(self, prefix=None):
+        common = ('..', '*') if prefix else ('*', )
+
+        if prefix is None:
+            prefix = self.app_prefix
+        d = self.apps
+        for p in prefix:
+            try:
+                d = d[p]
+            except:
+                return common
+        if isinstance(d, dict):
+            return common + tuple(sorted(d.keys()))
+        elif prefix:
+            return prefix[-1]
+        else:
+            return common + tuple(sorted(self.apps.keys()))
+
+
+__query_state = _QueryState()
+
+
+def _update_state(app_prefix,
+                  app_name,
+                  tags,
+                  offset,
+                  after,
+                  before,
+                  records_per_page,
+                  database='tcp://127.0.0.1:6789'):
+
+    state = __query_state
+
+    app_pattern = '.'.join(app_prefix + [app_name])
+    if app_pattern == '*':
+        app_pattern = None
+
+    if app_name == '..':
+        app_prefix = app_prefix[:-1]
+        app_name = '*'
+    elif app_name == '*':
+        pass
+    else:
+        app_prefix = app_prefix + [app_name]
+        app_name = '*'
+
+    if isinstance(database, str) and database.startswith('tcp://'):
+        with ZMQContextManager(zmq.DEALER, connect=database) as socket:
+            socket.send_pyobj({
+                'method': 'record_query',
+                'app': app_pattern,
+                'tags': tags,
+                'offset': offset,
+                'limit': records_per_page,
+                'before': before,
+                'after': after
+            })
+            total, apps, table = socket.recv_pyobj()
+    else:
+        from .curd import query_record
+        from .models import create_engine, sessionmaker
+
+        url = f'sqlite:///{database}/data.db'
+
+        engine = create_engine(url)
+        Session = sessionmaker(bind=engine)
+        with Session() as db:
+            total, apps, table = query_record(db,
+                                              offset=offset,
+                                              limit=records_per_page,
+                                              app=app_pattern,
+                                              tags=tags,
+                                              before=before,
+                                              after=after)
+    state.table = table
+    state.apps = apps
+    state.total = total
+    state.records_per_page = records_per_page
+    app_list = state.list_apps(app_prefix)
+    if isinstance(app_list, str):
+        state.app_name = app_prefix[-1]
+    else:
+        state.app_prefix = app_prefix
+        state.app_name = '*'
+    return total
+
+
+def _update_view(ui_widgets):
+    state = __query_state
+
+    page = state.page
+    total = state.total
+    records_per_page = state.records_per_page
+
+    app_options = state.list_apps(state.app_prefix)
+    ui_widgets['app_prefix'].value = 'App:' + '.'.join(state.app_prefix)
+    handlers = ui_widgets['app']._trait_notifiers['value']['change']
+    ui_widgets['app']._trait_notifiers['value']['change'] = handlers[:1]
+    ui_widgets['app'].options = app_options
+    ui_widgets['app']._trait_notifiers['value']['change'] = handlers
+    ui_widgets['app'].value = state.app_name
+
+    offset = (page - 1) * records_per_page
+    if offset < records_per_page:
+        ui_widgets['bt_pre'].disabled = True
+    else:
+        ui_widgets['bt_pre'].disabled = False
+
+    ui_widgets['table'].value = _format_table(state.table)
+    if total - offset <= records_per_page:
+        ui_widgets['bt_next'].disabled = True
+    else:
+        ui_widgets['bt_next'].disabled = False
+    ui_widgets[
+        'page_num_label'].value = f"{page} | {math.ceil(total/10)} pages"
+
+
+def _get_query_params(state: _QueryState, ui_widgets: dict):
+    page = state.page
+    records_per_page = state.records_per_page
+    tags = [t.strip() for t in ui_widgets['tags'].value.split(',') if t]
+    offset = (page - 1) * records_per_page
+    after = ui_widgets['after'].value
+    before = ui_widgets['before'].value
+
+    if state.app_name == '..' and len(state.app_prefix) == 1:
+        state.app_prefix = []
+        state.app_name = '*'
+
+    return state.app_prefix, state.app_name, tags, offset, after, before, records_per_page
+
+
+def _on_pre_bt_clicked(bt, ui_widgets):
+    __query_state.page -= 1
+    _update_state(*_get_query_params(__query_state, ui_widgets))
+    _update_view(ui_widgets)
+
+
+def _on_next_bt_clicked(bt, ui_widgets):
+    __query_state.page += 1
+    _update_state(*_get_query_params(__query_state, ui_widgets))
+    _update_view(ui_widgets)
+
+
+def _on_app_changed(changes, ui_widgets, stack=[]):
+    if changes['name'] not in ['value', 'index']:
+        stack.append(changes['name'])
+        return
+    if stack and changes['name'] == 'index':
+        stack.clear()
+        return
+
+    __query_state.app_name = ui_widgets['app'].value
+    __query_state.page = 1
+    _update_state(*_get_query_params(__query_state, ui_widgets))
+    _update_view(ui_widgets)
+
+
+def _on_after_changed(changes, ui_widgets):
+    if changes['name'] != 'value':
+        return
+    __query_state.page = 1
+    _update_state(*_get_query_params(__query_state, ui_widgets))
+    _update_view(ui_widgets)
+
+
+def _on_before_changed(changes, ui_widgets):
+    if changes['name'] != 'value':
+        return
+    __query_state.page = 1
+    _update_state(*_get_query_params(__query_state, ui_widgets))
+    _update_view(ui_widgets)
+
+
+def _on_tags_submit(tags, ui_widgets):
+    __query_state.page = 1
+    _update_state(*_get_query_params(__query_state, ui_widgets))
+    _update_view(ui_widgets)
+
+
+def lookup(app=None, limit=10, database=default_server):
+    after = widgets.DatePicker()
+    before = widgets.DatePicker()
+    app_prefix = widgets.Label('App:')
+    app_name = widgets.Dropdown(
+        options=[
+            '*',
+        ],
+        value='*',
+        disabled=False,
+    )
+    bt_pre = widgets.Button(description='<<')
+    bt_next = widgets.Button(description='>>')
+    tags = widgets.Text()
+    page_num_label = widgets.Label('1 | * pages')
+    row1 = widgets.HBox(
+        [bt_pre, page_num_label, bt_next, app_prefix, app_name])
+    row2 = widgets.HBox([widgets.Label('Between:'), after, before])
+    row3 = widgets.HBox([widgets.Label('  Tags :'), tags])
+    table = widgets.HTML()
+    box = widgets.VBox([row1, row2, row3, table])
+
+    ui_widgets = {
+        'app_prefix': app_prefix,
+        'app': app_name,
+        'before': before,
+        'after': after,
+        'tags': tags,
+        'table': table,
+        'page_num_label': page_num_label,
+        'bt_pre': bt_pre,
+        'bt_next': bt_next,
+    }
+
+    __query_state.records_per_page = limit
+    if app is not None:
+        __query_state.app_prefix = app.split('.')
+        __query_state.app_name = '*'
+
+    _update_state(*_get_query_params(__query_state, ui_widgets),
+                  database=database)
+    _update_view(ui_widgets)
+
+    bt_pre.on_click(
+        functools.partial(_on_pre_bt_clicked, ui_widgets=ui_widgets))
+    bt_next.on_click(
+        functools.partial(_on_next_bt_clicked, ui_widgets=ui_widgets))
+    after.observe(functools.partial(_on_after_changed, ui_widgets=ui_widgets),
+                  names='value')
+    before.observe(functools.partial(_on_before_changed,
+                                     ui_widgets=ui_widgets),
+                   names='value')
+    app_name.observe(functools.partial(_on_app_changed, ui_widgets=ui_widgets),
+                     names='value')
+    tags.on_submit(functools.partial(_on_tags_submit, ui_widgets=ui_widgets))
+
+    display(box)
+
+
+def lookup_list(*, full=False):
+    if __query_state.table is None:
+        return []
+    if full:
+        return __query_state.table['body']
+    else:
+        return [r[0] for r in __query_state.table['body']]
+
+
+def ping(database=default_server, timeout=1, socket=None):
+    with ZMQContextManager(zmq.DEALER,
+                           connect=database,
+                           socket=socket,
+                           timeout=timeout) as socket:
+
+        socket.send_pyobj({'method': 'ping'})
+        try:
+            return socket.recv_pyobj()
+        except zmq.Again:
+            raise TimeoutError(f"No response from server within {timeout} s")