nmreval/src/gui_qt/io/filereaders.py

from __future__ import annotations

from pathlib import Path
import struct

from ..Qt import QtCore
from .asciireader import QAsciiReader
from .hdfreader import QHdfViewer
from .bdsreader import QBDSReader
from .gracereader import QGraceReader
from .dscreader import QDSCReader
from .nmrreader import QNMRReader


class QFileReader(QtCore.QObject):
    data_read = QtCore.pyqtSignal([list], [dict])

    def __init__(self, manager=None):
        QtCore.QObject.__init__(self)

        self.select = 'all'
        self.data = []
        self.filenames = None
        self.extensions = set()

        self.reader = {}

        for ext, reader in [
            ('txt', QAsciiReader), ('dsc', QDSCReader), ('agr', QGraceReader),
            ('bds', QBDSReader), ('hdf', QHdfViewer),  ('nmr', QNMRReader)
        ]:
            self.register(ext, reader)

    def __call__(self, files: list[str] | str) -> list:
        self.data = []
        if isinstance(files, str):
            self.filenames = [files]
        else:
            self.filenames = files

        return self.readfiles(files)

    def readfiles(self, fname: list[str] | str) -> list:
        # clear old data
        self.data = []

        if not isinstance(fname, list):
            fname = [fname]

        for f in fname:
            f = Path(f)
            dtype = self.guess_type(f)
            if dtype in self.reader:
                r = self.reader[dtype]
            else:
                raise ValueError(f'Unknown type for file {f}')

            try:
                # If QAsciiReader.skip = True it accepts automatically and returns None
                r(f).exec()
            except ImportError:
                pass

        self.data_read.emit(self.data)

        self.reader['txt'].skip = False

        return self.data

    def readtnt(self, fname):
        """ Special treatment for tnt. If data is a single measurement, skip dialog """
        raise NotImplementedError
        # tntreader = self.reader[2]
        # tntreader(fname)
        # if not tntreader.reader.onedimensional:
        #     tntreader.exec()

    @QtCore.pyqtSlot(list)
    def _update_dic(self, other: list):
        self.data.extend(other)

    def register(self, key, new_reader):
        if key in self.reader:
            return self.reader[key]
        r = new_reader()
        self.reader[key] = r
        r.data_read.connect(self._update_dic)
        self.extensions.update(r.file_ext)

        return r

    def guess_type(self, fname: Path) -> str:
        ext = fname.suffix.lower()

        if ext in self.extensions:
            if ext in ('.dat', '.txt'):
                with fname.open('r', encoding='iso-8859-15') as fp:
                    line = fp.readline()
                    if line.strip().startswith('Filename:	C:\\'):
                        return 'dsc'

                return 'txt'

            if ext in ('.h5', '.hdf', '.hdf5'):
                return 'hdf'

            if ext == '.eps':
                return 'bds'

            return ext[1:]  # remove dot

        else:
            with fname.open('rb') as fp:
                s16 = fp.read(16)

                # copied from whichdb (Python 2.7) to look for magic value of dbhash
                (magic,) = struct.unpack("=l", s16[-4:])
                if magic in (0x00061561, 0x61150600):
                    return 'nmr'

                else:
                    ret_types = ('nmr', 'bds', 'tnt', 'agr', 'dsc')
                    strings = (b'NMREVAL', b'NOVOCONTROL', b'TNT1.005', b'# Grace project', b'Filename:\tC:')

                    (magic,) = struct.unpack('<16s', s16)
                    for dtype, startstring in zip(ret_types, strings):
                        if magic.startswith(startstring):
                            return dtype

                    # nothing matched, text file is best guess
                    return 'txt'