Merge pull request #14993 from jdavcs/dev_datatypes_contbl

Refactor, add type hints to get_chunk() in datatypes

lib/galaxy/datatypes/binary.py

@@ -13,7 +13,10 @@ import tarfile
 import tempfile
 import zipfile
 from json import dumps
-from typing import Optional
+from typing import (
+    Optional,
+    TYPE_CHECKING,
+)
 
 import h5py
 import numpy as np
@@ -56,6 +59,9 @@ from . import (
     dataproviders,
 )
 
+if TYPE_CHECKING:
+    from galaxy.model import DatasetInstance
+
 log = logging.getLogger(__name__)
 # pysam 0.16.0.1 emits logs containing the word 'Error', this can confuse the stdout/stderr checkers.
 # Can be be removed once https://github.com/pysam-developers/pysam/issues/939 is resolved.
@@ -571,7 +577,7 @@ class BamNative(CompressedArchive, _BamOrSam):
         # Remove temp file and empty temporary directory
         os.rmdir(tmp_dir)
 
-    def get_chunk(self, trans, dataset, offset=0, ck_size=None):
+    def get_chunk(self, trans, dataset: "DatasetInstance", offset: int = 0, ck_size: Optional[int] = None) -> str:
         if not offset == -1:
             try:
                 with pysam.AlignmentFile(dataset.file_name, "rb", check_sq=False) as bamfile:
@@ -579,8 +585,8 @@ class BamNative(CompressedArchive, _BamOrSam):
                     ck_data = ""
                     header_line_count = 0
                     if offset == 0:
-                        ck_data = bamfile.text.replace("\t", " ")
-                        header_line_count = bamfile.text.count("\n")
+                        ck_data = bamfile.text.replace("\t", " ")  # type: ignore[attr-defined]
+                        header_line_count = bamfile.text.count("\n")  # type: ignore[attr-defined]
                     else:
                         bamfile.seek(offset)
                     for line_number, alignment in enumerate(bamfile):

lib/galaxy/datatypes/tabular.py

@@ -12,6 +12,10 @@ import shutil
 import subprocess
 import tempfile
 from json import dumps
+from typing import (
+    Optional,
+    TYPE_CHECKING,
+)
 
 import pysam
 from markupsafe import escape
@@ -34,7 +38,6 @@ from galaxy.datatypes.sniff import (
     iter_headers,
     validate_tabular,
 )
-from galaxy.model import DatasetInstance
 from galaxy.util import compression_utils
 from galaxy.util.markdown import (
     indicate_data_truncated,
@@ -42,6 +45,9 @@ from galaxy.util.markdown import (
 )
 from . import dataproviders
 
+if TYPE_CHECKING:
+    from galaxy.model import DatasetInstance
+
 log = logging.getLogger(__name__)
 
 MAX_DATA_LINES = 100000
@@ -95,7 +101,7 @@ class TabularData(data.Text):
         if dataset.metadata.comment_lines:
             dataset.blurb = f"{dataset.blurb}, {util.commaify(str(dataset.metadata.comment_lines))} comments"
 
-    def displayable(self, dataset: DatasetInstance):
+    def displayable(self, dataset: "DatasetInstance"):
         try:
             return (
                 not dataset.dataset.purged
@@ -107,7 +113,17 @@ class TabularData(data.Text):
         except Exception:
             return False
 
-    def get_chunk(self, trans, dataset, offset=0, ck_size=None):
+    def get_chunk(self, trans, dataset: "DatasetInstance", offset: int = 0, ck_size: Optional[int] = None) -> str:
+        ck_data, last_read = self._read_chunk(trans, dataset, offset, ck_size)
+        return dumps(
+            {
+                "ck_data": util.unicodify(ck_data),
+                "offset": last_read,
+                "data_line_offset": self.data_line_offset,
+            }
+        )
+
+    def _read_chunk(self, trans, dataset: "DatasetInstance", offset: int, ck_size: Optional[int] = None):
         with compression_utils.get_fileobj(dataset.file_name) as f:
             f.seek(offset)
             ck_data = f.read(ck_size or trans.app.config.display_chunk_size)
@@ -117,13 +133,7 @@ class TabularData(data.Text):
                     ck_data += cursor
                     cursor = f.read(1)
             last_read = f.tell()
-        return dumps(
-            {
-                "ck_data": util.unicodify(ck_data),
-                "offset": last_read,
-                "data_line_offset": self.data_line_offset,
-            }
-        )
+        return ck_data, last_read
 
     def display_data(self, trans, dataset, preview=False, filename=None, to_ext=None, offset=None, ck_size=None, **kwd):
         headers = kwd.get("headers", {})
@@ -1541,29 +1551,26 @@ class ConnectivityTable(Tabular):
                 i += 1
         return False
 
-    def get_chunk(self, trans, dataset, chunk):
-        ck_index = int(chunk)
-        f = open(dataset.file_name)
-        f.seek(ck_index * trans.app.config.display_chunk_size)
-        # If we aren't at the start of the file, seek to next newline.  Do this better eventually.
-        if f.tell() != 0:
-            cursor = f.read(1)
-            while cursor and cursor != "\n":
-                cursor = f.read(1)
-        ck_data = f.read(trans.app.config.display_chunk_size)
-        cursor = f.read(1)
-        while cursor and ck_data[-1] != "\n":
-            ck_data += cursor
-            cursor = f.read(1)
+    def get_chunk(self, trans, dataset: "DatasetInstance", offset: int = 0, ck_size: Optional[int] = None) -> str:
+        ck_data, last_read = self._read_chunk(trans, dataset, offset, ck_size)
+        try:
+            # The ConnectivityTable format has several derivatives of which one is delimited by (multiple) spaces.
+            # By converting these spaces back to tabs, chunks can still be interpreted by tab delimited file parsers
+            ck_data_header, ck_data_body = ck_data.split("\n", 1)
+            ck_data_header = re.sub("^([0-9]+)[ ]+", r"\1\t", ck_data_header)
+            ck_data_body = re.sub("\n[ \t]+", "\n", ck_data_body)
+            ck_data_body = re.sub("[ ]+", "\t", ck_data_body)
+            ck_data = f"{ck_data_header}\n{ck_data_body}"
+        except ValueError:
+            pass  # 1 or 0 lines left
 
-        # The ConnectivityTable format has several derivatives of which one is delimited by (multiple) spaces.
-        # By converting these spaces back to tabs, chucks can still be interpreted by tab delimited file parsers
-        ck_data_header, ck_data_body = ck_data.split("\n", 1)
-        ck_data_header = re.sub("^([0-9]+)[ ]+", r"\1\t", ck_data_header)
-        ck_data_body = re.sub("\n[ \t]+", "\n", ck_data_body)
-        ck_data_body = re.sub("[ ]+", "\t", ck_data_body)
-
-        return dumps({"ck_data": util.unicodify(f"{ck_data_header}\n{ck_data_body}"), "ck_index": ck_index + 1})
+        return dumps(
+            {
+                "ck_data": util.unicodify(ck_data),
+                "offset": last_read,
+                "data_line_offset": self.data_line_offset,
+            }
+        )
 
 
 @build_sniff_from_prefix

test-data/1.ct

@@ -0,0 +1,3 @@
+363	tmRNA
+1	G	0	2	359	1
+2	G	1	3	358	2

test/unit/data/datatypes/test_connectivity_table.py

@@ -0,0 +1,49 @@
+import os
+
+import pytest
+
+from galaxy.datatypes.tabular import ConnectivityTable
+from galaxy.util import galaxy_directory
+
+
+@pytest.fixture
+def dataset():
+    class MockDataset:
+        file_name = os.path.join(galaxy_directory(), "test-data/1.ct")
+
+    return MockDataset()
+
+
+@pytest.fixture
+def make_trans():
+    class Stub:
+        pass
+
+    class StubTransaction:
+        def __init__(self, chunk_size):
+            self.app = Stub()
+            self.app.config = Stub()  # type: ignore[attr-defined]
+            self.app.config.display_chunk_size = chunk_size  # type: ignore[attr-defined]
+
+    return StubTransaction
+
+
+def test_get_chunk(dataset, make_trans):
+    dt = ConnectivityTable()
+    chunk_size = 1000
+    trans = make_trans(chunk_size)
+    chunk = dt.get_chunk(trans, dataset)
+    assert (
+        chunk
+        == '{"ck_data": "363\\ttmRNA\\n1\\tG\\t0\\t2\\t359\\t1\\n2\\tG\\t1\\t3\\t358\\t2\\n", "offset": 38, "data_line_offset": 0}'
+    )
+
+
+def test_get_chunk_with_offset(dataset, make_trans):
+    dt = ConnectivityTable()
+    chunk_size = 10
+    trans = make_trans(chunk_size)
+
+    # reads chunk_size chars from offset 5 (line 1) to the end of line 2
+    chunk = dt.get_chunk(trans, dataset, 5)
+    assert chunk == '{"ck_data": "mRNA\\n1\\tG\\t0\\t2\\t359\\t1", "offset": 24, "data_line_offset": 0}'