QuLab 2.1.1__tar.gz → 2.1.3__tar.gz

{qulab-2.1.1 → qulab-2.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: QuLab
-Version: 2.1.1
+Version: 2.1.3
 Summary: contral instruments and manage data
 Author-email: feihoo87 <feihoo87@gmail.com>
 Maintainer-email: feihoo87 <feihoo87@gmail.com>

{qulab-2.1.1 → qulab-2.1.3}/QuLab.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: QuLab
-Version: 2.1.1
+Version: 2.1.3
 Summary: contral instruments and manage data
 Author-email: feihoo87 <feihoo87@gmail.com>
 Maintainer-email: feihoo87 <feihoo87@gmail.com>

{qulab-2.1.1 → qulab-2.1.3}/QuLab.egg-info/SOURCES.txt

@@ -29,7 +29,6 @@ qulab/scan/models.py
 qulab/scan/optimize.py
 qulab/scan/query.py
 qulab/scan/record.py
-qulab/scan/recorder.py
 qulab/scan/scan.py
 qulab/scan/server.py
 qulab/scan/space.py

qulab-2.1.3/qulab/__init__.py

@@ -0,0 +1,3 @@
+from .scan import Scan, get_record, load_record, lookup, lookup_list
+from .version import __version__
+from .visualization import autoplot

{qulab-2.1.1 → qulab-2.1.3}/qulab/__main__.py

@@ -1,7 +1,6 @@
 import click
 
 from .monitor.__main__ import main as monitor
-from .scan.recorder import record
 from .scan.server import server
 from .sys.net.cli import dht
 from .visualization.__main__ import plot
@@ -21,7 +20,6 @@ def hello():
 main.add_command(monitor)
 main.add_command(plot)
 main.add_command(dht)
-main.add_command(record)
 main.add_command(server)
 
 if __name__ == '__main__':

{qulab-2.1.1 → qulab-2.1.3}/qulab/scan/__init__.py

@@ -1,3 +1,3 @@
 from .expression import Expression, Symbol
-from .query import get_record, lookup, lookup_list
+from .query import get_record, load_record, lookup, lookup_list
 from .scan import Scan

{qulab-2.1.1 → qulab-2.1.3}/qulab/scan/curd.py

@@ -1,4 +1,7 @@
+import lzma
+import pickle
 from datetime import date, datetime, timezone
+from pathlib import Path
 from typing import Sequence, Type, Union
 
 from sqlalchemy.orm import Query, Session, aliased
@@ -6,7 +9,8 @@ from sqlalchemy.orm.exc import NoResultFound
 from sqlalchemy.orm.session import Session
 from waveforms.dicttree import foldDict
 
-from .models import Comment, Record, Report, Sample, Tag
+from .models import (Cell, Comment, Config, InputText, Notebook, Record,
+                     Report, Sample, Tag, utcnow)
 
 
 def tag(session: Session, tag_text: str) -> Tag:
@@ -142,3 +146,76 @@ def remove_tags(session: Session, record_id: int, tags: Sequence[str]):
         session.rollback()
         return False
     return True
+
+
+def create_notebook(session: Session, notebook_name: str) -> Notebook:
+    """Create a notebook in the database."""
+    notebook = Notebook(name=notebook_name)
+    session.add(notebook)
+    return notebook
+
+
+def create_input_text(session: Session, input_text: str) -> InputText:
+    """Create an input text in the database."""
+    input = InputText()
+    input.text = input_text
+    try:
+        input = session.query(InputText).filter(
+            InputText.hash == input.hash,
+            InputText.text_field == input_text).one()
+    except NoResultFound:
+        session.add(input)
+    return input
+
+
+def create_cell(session: Session, notebook: Notebook, input_text: str) -> Cell:
+    """Create a cell in the database."""
+    cell = Cell()
+    cell.notebook = notebook
+    cell.input = create_input_text(session, input_text)
+    cell.index = len(notebook.cells) - 1
+    session.add(cell)
+    notebook.atime = cell.ctime
+    return cell
+
+
+def create_config(session: Session, config: dict | bytes, base: Path,
+                  filename: str) -> Config:
+    """Create a config in the database."""
+
+    if not isinstance(config, bytes):
+        buf = pickle.dumps(config)
+        buf = lzma.compress(buf)
+        content_type = 'application/pickle+lzma'
+    else:
+        buf = config
+        content_type = 'application/octet-stream'
+    config = Config(buf)
+    config.content_type = content_type
+    for cfg in session.query(Config).filter(Config.hash == config.hash).all():
+        with open(base / cfg.file, 'rb') as f:
+            if f.read() == buf:
+                cfg.atime = utcnow()
+                return cfg
+    else:
+        path = base / filename
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with open(path, 'wb') as f:
+            f.write(buf)
+        config.file = filename
+        session.add(config)
+    return config
+
+
+def get_config(session: Session, config_id: int, base: Path):
+    config = session.get(Config, config_id)
+    if config is None:
+        return None
+    config.atime = utcnow()
+    path = base / config.file
+    with open(path, 'rb') as f:
+        buf = f.read()
+    if config.content_type == 'application/pickle+lzma':
+        buf = lzma.decompress(buf)
+        buf = pickle.loads(buf)
+    return buf

{qulab-2.1.1 → qulab-2.1.3}/qulab/scan/models.py

@@ -1,11 +1,9 @@
 import hashlib
 import pickle
-import time
 from datetime import datetime, timezone
 from functools import singledispatchmethod
-from typing import Optional
 
-from sqlalchemy import (JSON, Column, DateTime, Float, ForeignKey, Integer,
+from sqlalchemy import (Column, DateTime, Float, ForeignKey, Integer,
                         LargeBinary, String, Table, Text, create_engine)
 from sqlalchemy.orm import (backref, declarative_base, relationship,
                             sessionmaker)
@@ -325,7 +323,7 @@ class InputText(Base):
     __tablename__ = 'inputs'
 
     id = Column(Integer, primary_key=True)
-    hash = Column(LargeBinary(20))
+    hash = Column(LargeBinary(20), index=True)
     text_field = Column(Text, unique=True)
 
     @property
@@ -432,6 +430,22 @@ class SampleTransfer(Base):
     comments = relationship("Comment", secondary=sample_transfer_comments)
 
 
+class Config(Base):
+    __tablename__ = 'configs'
+
+    id = Column(Integer, primary_key=True)
+    hash = Column(LargeBinary(20), index=True)
+    file = Column(String)
+    content_type = Column(String, default='application/pickle')
+    ctime = Column(DateTime, default=utcnow)
+    atime = Column(DateTime, default=utcnow)
+
+    records = relationship("Record", back_populates="config")
+
+    def __init__(self, data: bytes) -> None:
+        self.hash = hashlib.sha1(data).digest()
+
+
 class Record(Base):
     __tablename__ = 'records'
 
@@ -440,14 +454,14 @@ class Record(Base):
     mtime = Column(DateTime, default=utcnow)
     atime = Column(DateTime, default=utcnow)
     user_id = Column(Integer, ForeignKey('users.id'))
+    config_id = Column(Integer, ForeignKey('configs.id'))
     parent_id = Column(Integer, ForeignKey('records.id'))
     cell_id = Column(Integer, ForeignKey('cells.id'))
 
     app = Column(String)
     file = Column(String)
+    content_type = Column(String, default='application/pickle')
     key = Column(String)
-    config = Column(JSON)
-    task_hash = Column(LargeBinary(32))
 
     parent = relationship("Record",
                           remote_side=[id],
@@ -456,6 +470,7 @@ class Record(Base):
                             remote_side=[parent_id],
                             back_populates="parent")
 
+    config = relationship("Config", back_populates="records")
     user = relationship("User")
     samples = relationship("Sample",
                            secondary=sample_records,

{qulab-2.1.1 → qulab-2.1.3}/qulab/scan/query.py

@@ -6,10 +6,13 @@ 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 .server import get_local_record
 from .scan import default_server
 
 
@@ -26,18 +29,15 @@ def get_record(id, database=default_server) -> Record:
             d._file = None
             return d
     else:
-        from .models import Record as RecordInDB
-        from .models import create_engine, sessionmaker
-
         db_file = Path(database) / 'data.db'
         engine = create_engine(f'sqlite:///{db_file}')
         Session = sessionmaker(bind=engine)
         with Session() as session:
-            path = Path(database) / 'objects' / session.get(RecordInDB,
-                                                            id).file
-            with open(path, 'rb') as f:
-                record = dill.load(f)
-            return record
+            return get_local_record(session, id, database)
+
+
+def load_record(file):
+    return Record.load(file)
 
 
 def _format_tag(tag):

{qulab-2.1.1 → qulab-2.1.3}/qulab/scan/record.py

@@ -1,6 +1,7 @@
 import itertools
 import sys
 import uuid
+import zipfile
 from pathlib import Path
 from threading import Lock
 from types import EllipsisType
@@ -13,7 +14,7 @@ from qulab.sys.rpc.zmq_socket import ZMQContextManager
 
 from .space import OptimizeSpace, Space
 
-_notgiven = object()
+_not_given = object()
 
 
 def random_path(base):
@@ -94,6 +95,11 @@ class BufferList():
                         dill.dump(item, f)
                 self._list.clear()
 
+    def delete(self):
+        if isinstance(self.file, Path):
+            self.file.unlink()
+            self.file = None
+
     def append(self, pos, value, dims=None):
         if dims is not None:
             if any([p != 0 for i, p in enumerate(pos) if i not in dims]):
@@ -120,6 +126,18 @@ class BufferList():
                             yield pos, value
                         except EOFError:
                             break
+        elif isinstance(
+                self.file, tuple) and len(self.file) == 2 and isinstance(
+                    self.file[0], str) and self.file[0].endswith('.zip'):
+            f, name = self.file
+            with zipfile.ZipFile(f, 'r') as z:
+                with z.open(name, 'r') as f:
+                    while True:
+                        try:
+                            pos, value = dill.load(f)
+                            yield pos, value
+                        except EOFError:
+                            break
 
     def iter(self):
         if self._data_id is None:
@@ -327,7 +345,7 @@ class Record():
             ret = ret.toarray()
         return ret
 
-    def get(self, key, default=_notgiven, buffer_to_array=False, slice=None):
+    def get(self, key, default=_not_given, buffer_to_array=False, slice=None):
         if self.is_remote_record():
             with ZMQContextManager(zmq.DEALER,
                                    connect=self.database) as socket:
@@ -355,7 +373,7 @@ class Record():
                 else:
                     return ret
         else:
-            if default is _notgiven:
+            if default is _not_given:
                 d = self._items.get(key)
             else:
                 d = self._items.get(key, default)
@@ -437,6 +455,79 @@ class Record():
         with open(self._file, 'wb') as f:
             dill.dump(self, f)
 
+    def delete(self):
+        if self.is_remote_record():
+            with ZMQContextManager(zmq.DEALER,
+                                   connect=self.database) as socket:
+                socket.send_pyobj({
+                    'method': 'record_delete',
+                    'record_id': self.id
+                })
+        elif self.is_local_record():
+            for key, value in self._items.items():
+                if isinstance(value, BufferList):
+                    value.delete()
+            self._file.unlink()
+
+    def export(self, file):
+        with zipfile.ZipFile(file,
+                             'w',
+                             compression=zipfile.ZIP_DEFLATED,
+                             compresslevel=9) as z:
+            items = {}
+            for key in self.keys():
+                value = self.get(key)
+                if isinstance(value, BufferList):
+                    v = BufferList()
+                    v.lu = value.lu
+                    v.rd = value.rd
+                    v.inner_shape = value.inner_shape
+                    items[key] = v
+                    with z.open(f'{key}.buf', 'w') as f:
+                        for pos, data in value.iter():
+                            dill.dump((pos, data), f)
+                else:
+                    items[key] = value
+            with z.open('record.pkl', 'w') as f:
+                self.description['entry']['scripts'] = self.scripts()
+                dill.dump((self.description, items), f)
+
+    def scripts(self, session=None):
+        scripts = self.description['entry']['scripts']
+        if isinstance(scripts, list):
+            return scripts
+        else:
+            cell_id = scripts
+
+        if self.is_remote_record():
+            with ZMQContextManager(zmq.DEALER,
+                                   connect=self.database) as socket:
+                socket.send_pyobj({
+                    'method': 'notebook_history',
+                    'cell_id': cell_id
+                })
+                return socket.recv_pyobj()
+        elif self.is_local_record():
+            from .models import Cell
+            assert session is not None, "session is required for local record"
+            cell = session.get(Cell, cell_id)
+            return [
+                cell.input.text
+                for cell in cell.notebook.cells[1:cell.index + 2]
+            ]
+
+    @classmethod
+    def load(cls, file: str):
+        with zipfile.ZipFile(file, 'r') as z:
+            with z.open('record.pkl', 'r') as f:
+                description, items = dill.load(f)
+            record = cls(None, None, description)
+            for key, value in items.items():
+                if isinstance(value, BufferList):
+                    value.file = file, f'{key}.buf'
+                record._items[key] = value
+        return record
+
     def __repr__(self):
         return f"<Record: id={self.id} app={self.description['app']}, keys={self.keys()}>"