QuLab 2.10.10__cp313-cp313-win_amd64.whl

qulab/storage/file.py

@@ -0,0 +1,273 @@
+import lzma
+import pathlib
+
+import dill
+import numpy as np
+
+MAGIC = b'wAvEDatA'
+
+lzma_filters = [{
+    'id': lzma.FILTER_LZMA2,
+    'preset': 9 | lzma.PRESET_EXTREME,
+}]
+
+FORMAT_RAW = 1
+FORMAT_XZ = 2
+
+
+def _get_header(file):
+    assert MAGIC == file.read(len(MAGIC)), 'Invalid file format'
+    version = file.read(1)
+    version = int.from_bytes(version, 'big', signed=False)
+    file_format = file.read(1)
+    file_format = int.from_bytes(file_format, 'big', signed=False)
+    info_size = file.read(2)
+    info_size = int.from_bytes(info_size, 'big', signed=False)
+    header = lzma.decompress(file.read(info_size),
+                             format=lzma.FORMAT_RAW,
+                             filters=lzma_filters)
+    info = dill.loads(header)
+    return {'version': version, 'format': file_format, 'info': info}
+
+
+def _make_header(header: dict):
+    version = header.get('version', 1)
+    file_format = header.get('format', 1)
+    info = header.get('info', {})
+    header = lzma.compress(dill.dumps(info),
+                           format=lzma.FORMAT_RAW,
+                           filters=lzma_filters)
+    assert len(header) < 2**16, 'Header too large'
+    version = version.to_bytes(1, 'big', signed=False)
+    file_format = file_format.to_bytes(1, 'big', signed=False)
+    info_size = len(header).to_bytes(2, 'big', signed=False)
+    return MAGIC + version + file_format + info_size + header
+
+
+class FakeLock():
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, *args):
+        pass
+
+    async def __aenter__(self):
+        pass
+
+    async def __aexit__(self, *args):
+        pass
+
+
+class BaseFile():
+
+    def __init__(self, path: str | pathlib.Path, lock=None):
+        self.path = path
+        self.lock = lock or FakeLock()
+
+    def compress(self):
+        with self.lock:
+            self._compress()
+
+    def decompress(self):
+        with self.lock:
+            self._decompress()
+
+    def _compress(self):
+        with open(self.path, 'rb+') as f:
+            try:
+                header = _get_header(f)
+            except:
+                return
+            if header.get('format', FORMAT_RAW) == FORMAT_XZ:
+                return
+            header['format'] = FORMAT_XZ
+            buffer = f.read()
+            f.seek(0)
+            f.write(_make_header(header))
+            f.write(
+                lzma.compress(buffer,
+                              format=lzma.FORMAT_RAW,
+                              filters=lzma_filters))
+            f.truncate()
+
+    def _decompress(self):
+        with open(self.path, 'rb+') as f:
+            try:
+                header = _get_header(f)
+            except:
+                return
+            if header.get('format', FORMAT_RAW) == FORMAT_RAW:
+                return
+            header['format'] = FORMAT_RAW
+            buffer = lzma.decompress(f.read(),
+                                     format=lzma.FORMAT_RAW,
+                                     filters=lzma_filters)
+            f.seek(0)
+            f.write(_make_header(header))
+            f.write(buffer)
+            f.truncate()
+
+
+class BinaryFile(BaseFile):
+
+    def read(self):
+        with self.lock:
+            with open(self.path, 'rb') as f:
+                header = _get_header(f)
+                if header.get('format', FORMAT_RAW) == FORMAT_XZ:
+                    return lzma.decompress(f.read(),
+                                           format=lzma.FORMAT_RAW,
+                                           filters=lzma_filters)
+                else:
+                    return f.read()
+
+    def write(self, data: bytes):
+        with self.lock:
+            with open(self.path, 'wb') as f:
+                f.write(
+                    _make_header({
+                        'format': FORMAT_RAW,
+                        'info': {
+                            'type': 'bytes'
+                        }
+                    }))
+                f.write(data)
+
+
+class ObjectFile(BinaryFile):
+
+    def load(self):
+        with self.lock:
+            with open(self.path, 'rb') as f:
+                header = _get_header(f)
+                if header.get('format', FORMAT_RAW) == FORMAT_XZ:
+                    return dill.loads(
+                        lzma.decompress(f.read(),
+                                        format=lzma.FORMAT_RAW,
+                                        filters=lzma_filters))
+                else:
+                    return dill.load(f)
+
+    def dump(self, obj):
+        with self.lock:
+            with open(self.path, 'wb') as f:
+                f.write(
+                    _make_header({
+                        'format': FORMAT_RAW,
+                        'info': {
+                            'type': 'pickle'
+                        }
+                    }))
+                dill.dump(obj, f)
+
+
+class ObjectListFile(ObjectFile):
+
+    def clear(self):
+        with self.lock:
+            self._decompress()
+            with open(self.path, 'wb') as f:
+                f.write(
+                    _make_header({
+                        'format': FORMAT_RAW,
+                        'info': {
+                            'type': 'pickle_list'
+                        }
+                    }))
+
+    def append(self, obj):
+        with self.lock:
+            with open(self.path, 'ab') as f:
+                if f.tell() == 0:
+                    f.write(
+                        _make_header({
+                            'format': FORMAT_RAW,
+                            'info': {
+                                'type': 'pickle_list'
+                            }
+                        }))
+                dill.dump(obj, f)
+
+    def __iter__(self):
+        with self.lock:
+            with open(self.path, 'rb') as f:
+                while True:
+                    try:
+                        header = _get_header(f)
+                    except AssertionError:
+                        break
+                    if header.get('format', FORMAT_RAW) == FORMAT_XZ:
+                        yield from self._compressed_iter(f)
+                    else:
+                        yield from self._iter(f)
+
+    def _compressed_iter(self, f):
+        with lzma.open(f, 'rb', format=lzma.FORMAT_RAW,
+                       filters=lzma_filters) as f:
+            yield from self._iter(f)
+
+    def _iter(self, f):
+        while True:
+            try:
+                yield dill.load(f)
+            except EOFError:
+                break
+
+    def asarray(self):
+        return np.array(list(self))
+
+
+class ArrayFile(BaseFile):
+
+    def clear(self):
+        with self.lock:
+            with open(self.path, 'wb') as f:
+                f.write(b'')
+
+    def extend(self, data: np.ndarray):
+        with self.lock:
+            self._decompress()
+            with open(self.path, 'ab') as f:
+                if f.tell() == 0:
+                    f.write(
+                        _make_header({
+                            'format': FORMAT_RAW,
+                            'info': {
+                                'type': 'array',
+                                'dtype': data.dtype
+                            }
+                        }))
+                f.write(data.tobytes())
+
+    def append(self, data, dtype=float):
+        self.extend(np.asarray(data, dtype=dtype))
+
+    def asarray(self):
+        with self.lock:
+            with open(self.path, 'rb') as f:
+                header = _get_header(f)
+                dtype = header['info']['dtype']
+                if header.get('format', FORMAT_RAW) == FORMAT_XZ:
+                    return np.frombuffer(
+                        lzma.decompress(f.read(),
+                                        format=lzma.FORMAT_RAW,
+                                        filters=lzma_filters),
+                        dtype=dtype)
+                else:
+                    return np.fromfile(f, dtype=dtype)
+
+
+def load(path):
+    with open(path, 'rb') as f:
+        header = _get_header(f)
+        if header['info']['type'] == 'array':
+            return ArrayFile(path)
+        elif header['info']['type'] == 'pickle':
+            return ObjectFile(path)
+        elif header['info']['type'] == 'pickle_list':
+            return ObjectListFile(path)
+        elif header['info']['type'] == 'bytes':
+            return BinaryFile(path)
+        else:
+            raise ValueError(f'Unknown file type: {header["info"]["type"]}')

qulab/storage/models/__init__.py

@@ -0,0 +1,22 @@
+from sqlalchemy import Engine
+from sqlalchemy.orm import Session
+
+from .base import Base
+from .file import File, FileChunk
+from .ipy import Cell, InputText, Notebook
+from .record import Record
+from .tag import Tag, has_tags
+
+
+def create_tables(url: str):
+    from sqlalchemy import create_engine
+    engine = create_engine(url)
+    Base.metadata.create_all(engine)
+
+
+def create_session(url: str) -> Session:
+    from sqlalchemy import create_engine
+    from sqlalchemy.orm import sessionmaker
+    engine = create_engine(url)
+    session = sessionmaker(bind=engine)
+    return session()

qulab/storage/models/base.py

@@ -0,0 +1,4 @@
+from sqlalchemy.orm import declarative_base
+from sqlalchemy.orm.session import Session
+
+Base = declarative_base()

qulab/storage/models/config.py

@@ -0,0 +1,28 @@
+import dill
+from sqlalchemy import Column, ForeignKey, Integer, String
+from sqlalchemy.orm import relationship
+
+from .base import Base
+from .file import File
+from .tag import Tag, has_tags
+
+
+@has_tags
+class Config(Base):
+    __tablename__ = 'configs'
+
+    id = Column(Integer, primary_key=True)
+    name = Column(String)
+    file_id = Column(Integer, ForeignKey('files.id'))
+    meta_id = Column(Integer, ForeignKey('files.id'))
+
+    file = relationship("File", foreign_keys=[file_id])
+
+    @property
+    def data(self) -> dict:
+        result = dill.loads(self.file.read())
+        return result
+
+    @data.setter
+    def data(self, data: dict):
+        self.file.write(dill.dumps(data))

qulab/storage/models/file.py

@@ -0,0 +1,89 @@
+import zlib
+from datetime import datetime
+
+from sqlalchemy import (Boolean, Column, DateTime, ForeignKey, Integer, String,
+                        event)
+from sqlalchemy.orm import relationship
+from sqlalchemy.orm.session import Session
+
+from ..chunk import CHUNKSIZE, delete_chunk, load_chunk, save_chunk
+from . import Base
+
+
+class FileChunk(Base):
+    __tablename__ = 'file_chunks'
+
+    id = Column(Integer, primary_key=True)
+    file_id = Column(Integer, ForeignKey('files.id'))
+    index = Column(Integer)
+    size = Column(Integer)
+    chunk_hash = Column(String)
+    compressed = Column(Boolean, default=False)
+
+    file = relationship("File", back_populates="chunks")
+
+    @property
+    def chunk(self):
+        if hasattr(self, '_chunk'):
+            return self._chunk
+        return load_chunk(self.chunk_hash, self.compressed)
+
+    @chunk.setter
+    def chunk(self, data):
+        self._chunk = data
+
+
+@event.listens_for(FileChunk, 'before_insert')
+def before_insert_file_chunk(mapper, connection, target: FileChunk):
+    target.chunk_hash, target.size = save_chunk(target._chunk)
+
+
+class File(Base):
+    __tablename__ = 'files'
+
+    id = Column(Integer, primary_key=True)
+    ctime = Column(DateTime, default=datetime.utcnow)
+    mtime = Column(DateTime, default=datetime.utcnow)
+    atime = Column(DateTime, default=datetime.utcnow)
+    name = Column(String)
+    size = Column(Integer)
+    chunks = relationship("FileChunk", order_by=FileChunk.index)
+
+    def write(self, data):
+        self.mtime = datetime.utcnow()
+        self.size = len(data)
+        self.chunks = []
+        for i in range(0, len(data), CHUNKSIZE):
+            chunk = FileChunk()
+            chunk.index = i // CHUNKSIZE
+            chunk._chunk = data[i:i + CHUNKSIZE]
+            chunk.size = len(chunk._chunk)
+            self.chunks.append(chunk)
+
+    def read(self):
+        self.atime = datetime.utcnow()
+        return b''.join([c.chunk for c in self.chunks])
+
+
+@event.listens_for(File, 'before_delete')
+def before_delete_file(mapper, connection, target):
+    for chunk in target.chunks:
+        connection.execute(
+            FileChunk.__table__.delete().where(FileChunk.id == chunk.id))
+
+
+def compress_chunks(db: Session):
+    for chunk in db.query(FileChunk).filter(
+            FileChunk.compressed == False).limit(100):
+        old_chunk_hash = chunk.chunk_hash
+        buf = zlib.compress(chunk.chunk)
+        chunk_hash, size = save_chunk(buf, compressed=True)
+        chunk.chunk_hash = chunk_hash
+        chunk.size = size
+        chunk.compressed = True
+        try:
+            db.commit()
+            delete_chunk(old_chunk_hash)
+        except:
+            db.rollback()
+            delete_chunk(chunk_hash)

qulab/storage/models/ipy.py

@@ -0,0 +1,58 @@
+import hashlib
+from datetime import datetime
+
+from sqlalchemy import (Column, DateTime, ForeignKey, Integer, LargeBinary,
+                        String, Text)
+from sqlalchemy.orm import relationship
+
+from .base import Base
+
+
+class InputText(Base):
+    __tablename__ = 'inputs'
+
+    id = Column(Integer, primary_key=True)
+    hash = Column(LargeBinary(20))
+    text_field = Column(Text, unique=True)
+
+    @property
+    def text(self):
+        return self.text_field
+
+    @text.setter
+    def text(self, text):
+        self.hash = hashlib.sha1(text.encode('utf-8')).digest()
+        self.text_field = text
+
+    def __repr__(self) -> str:
+        return self.text
+
+
+class Cell(Base):
+    __tablename__ = 'cells'
+
+    id = Column(Integer, primary_key=True)
+    notebook_id = Column(Integer, ForeignKey("notebooks.id"))
+    index = Column(Integer)
+    ctime = Column(DateTime, default=datetime.utcnow)
+    ftime = Column(DateTime, default=datetime.utcnow)
+    input_id = Column(Integer, ForeignKey("inputs.id"))
+
+    notebook = relationship("Notebook", back_populates="cells")
+    input = relationship("InputText")
+
+    def __repr__(self) -> str:
+        return f"Cell(index={self.index}, input='{self.input}')"
+
+
+class Notebook(Base):
+    __tablename__ = 'notebooks'
+
+    id = Column(Integer, primary_key=True)
+    name = Column(String)
+    ctime = Column(DateTime, default=datetime.utcnow)
+    atime = Column(DateTime, default=datetime.utcnow)
+
+    cells = relationship("Cell",
+                         order_by=Cell.index,
+                         back_populates="notebook")

qulab/storage/models/models.py

@@ -0,0 +1,88 @@
+from datetime import datetime
+
+from sqlalchemy import (Column, DateTime, Float, ForeignKey, Integer, String,
+                        Table, create_engine)
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import relationship, sessionmaker
+
+Base = declarative_base()
+
+dataset_tag_association = Table(
+    'dataset_tag', Base.metadata,
+    Column('dataset_id', Integer, ForeignKey('datasets.id'), primary_key=True),
+    Column('tag_id', Integer, ForeignKey('tags.id'), primary_key=True))
+
+
+class Experiment(Base):
+    __tablename__ = 'experiments'
+    id = Column(Integer, primary_key=True)
+    name = Column(String)
+    description = Column(String)
+    start_date = Column(DateTime, default=datetime.utcnow)
+    end_date = Column(DateTime, default=datetime.utcnow)
+    researcher_id = Column(Integer, ForeignKey('researchers.id'))
+    datasets = relationship('Dataset', back_populates='experiment')
+
+
+class Researcher(Base):
+    __tablename__ = 'researchers'
+    id = Column(Integer, primary_key=True)
+    first_name = Column(String)
+    last_name = Column(String)
+    email = Column(String)
+    affiliation = Column(String)
+    experiments = relationship('Experiment', back_populates='researcher')
+
+
+class Dataset(Base):
+    __tablename__ = 'datasets'
+    id = Column(Integer, primary_key=True)
+    name = Column(String)
+    description = Column(String)
+    tags = relationship('Tag',
+                        secondary=dataset_tag_association,
+                        back_populates='datasets')
+    creation_date = Column(DateTime, default=datetime.utcnow)
+    modification_date = Column(DateTime, default=datetime.utcnow)
+    experiment_id = Column(Integer, ForeignKey('experiments.id'))
+    experiment = relationship('Experiment', back_populates='datasets')
+    data_entries = relationship('DataEntry', back_populates='dataset')
+
+
+class DataEntry(Base):
+    __tablename__ = 'data_entries'
+    id = Column(Integer, primary_key=True)
+    timestamp = Column(DateTime)
+    value = Column(Float)
+    units = Column(String)
+    variable_name = Column(String)
+    dataset_id = Column(Integer, ForeignKey('datasets.id'))
+    dataset = relationship('Dataset', back_populates='data_entries')
+
+
+class Variable(Base):
+    __tablename__ = 'variables'
+    id = Column(Integer, primary_key=True)
+    name = Column(String)
+    description = Column(String)
+    data_type = Column(String)
+
+
+class Tag(Base):
+    __tablename__ = 'tags'
+    id = Column(Integer, primary_key=True)
+    name = Column(String, unique=True)
+    datasets = relationship('Dataset',
+                            secondary=dataset_tag_association,
+                            back_populates='tags')
+
+
+# 创建数据库引擎
+engine = create_engine('sqlite:///experiment_data.db')
+
+# 创建表
+Base.metadata.create_all(engine)
+
+# 创建会话
+Session = sessionmaker(bind=engine)
+session = Session()