Add AnnData as format

api/python/quilt3/formats.py

@@ -318,6 +318,7 @@ def all_supported_formats(cls):
 
                             Python Object Type     Serialization Formats
          <class 'pandas.core.frame.DataFrame'>  [ssv, csv, tsv, parquet]
+                     <class 'anndata.AnnData'>                   [.h5ad]
                        <class 'numpy.ndarray'>                [npy, npz]
                                  <class 'str'>      [md, json, rst, txt]
                                 <class 'dict'>                    [json]
@@ -342,6 +343,13 @@ def all_supported_formats(cls):
         else:
             cls.search(pd.DataFrame)  # Force FormatHandlers to register pd.DataFrame as a supported object type
 
+        try:
+            import anndata as ad
+        except ImportError:
+            pass
+        else:
+            cls.search(ad.AnnData)  # Force FormatHandlers to register ad.AnnData as a supported object type
+
         type_map = defaultdict(set)
         for handler in cls.registered_handlers:
             for t in handler.handled_types:
@@ -1024,6 +1032,56 @@ def deserialize(self, bytes_obj, meta=None, ext=None, **format_opts):
 ParquetFormatHandler().register()  # latest is preferred
 
 
+# noinspection PyPackageRequirements
+class AnnDataFormatHandler(BaseFormatHandler):
+    """Format for AnnData <--> .h5ad
+
+    Format Opts:
+        The following options may be used anywhere format opts are accepted,
+        or directly in metadata under `{'format': {'opts': {...: ...}}}`.
+
+        compression('gzip', 'lzf', None):  applies during serialization only.
+
+    """
+    name = 'h5ad'
+    handled_extensions = ['h5ad']
+    opts = ('compression',)
+    defaults = dict(
+        compression='lzf',
+    )
+
+    def handles_type(self, typ: type) -> bool:
+        # don't load module unless we actually have to use it.
+        if 'annndata' not in sys.modules:
+            return False
+        import anndata as ad
+        self.handled_types.add(ad.AnnData)
+        return super().handles_type(typ)
+
+    def serialize(self, obj, meta=None, ext=None, **format_opts):
+        import anndata as ad
+
+        opts = self.get_opts(meta, format_opts)
+        opts_with_defaults = copy.deepcopy(self.defaults)
+        opts_with_defaults.update(opts)
+        buf = io.BytesIO()
+        obj.write(buf, **opts_with_defaults)
+
+        return buf.getvalue(), self._update_meta(meta, additions=opts_with_defaults)
+
+    def deserialize(self, bytes_obj, meta=None, ext=None, **format_opts):
+        try:
+            import anndata as ad
+        except ImportError:
+            raise QuiltException("Please install anndata")
+
+        buf = io.BytesIO(bytes_obj)
+        return ad.read_h5ad(buf)
+
+
+AnnDataFormatHandler().register()
+
+
 class CompressionRegistry:
     """A collection for organizing `CompressionHandler` objects."""
     registered_handlers = []

api/python/setup.py

@@ -73,10 +73,9 @@ def run(self):
             'pandas>=0.19.2',
             'pyarrow>=0.14.1',              # as of 7/5/19: linux/circleci bugs on 0.14.0
         ],
+        'anndata': ['anndata>=0.8.0'],
         'tests': [
-            'numpy>=1.14.0',                # required by pandas, but missing from its dependencies.
-            'pandas>=0.19.2',
-            'pyarrow>=0.14.1',              # as of 7/5/19: linux/circleci bugs on 0.14.0
+            'quilt3[pyarrow,anndata]',
             'pytest==6.*',
             'pytest-cov',
             'coverage==6.4',

api/python/tests/test_formats.py

@@ -3,30 +3,33 @@
 import numpy as np
 import pandas as pd
 import pytest
+from anndata import AnnData
 
 from quilt3.formats import FormatRegistry
 from quilt3.util import QuiltException
 
 # Constants
 
+data_dir = pathlib.Path(__file__).parent / 'data'
+
 
 # Code
-def test_buggy_parquet():
+@pytest.mark.parametrize('parquet_handler', FormatRegistry.for_format('parquet'))
+def test_buggy_parquet(parquet_handler):
     """
     Test that Quilt avoids crashing on bad Pandas metadata from
     old pyarrow libaries.
     """
-    path = pathlib.Path(__file__).parent
-    for parquet_handler in FormatRegistry.for_format('parquet'):
-        with open(path / 'data' / 'buggy_parquet.parquet', 'rb') as bad_parq:
-            # Make sure this doesn't crash.
-            parquet_handler.deserialize(bad_parq.read())
+    path = data_dir / 'buggy_parquet.parquet'
+    data = path.read_bytes()
+    # Make sure this doesn't crash.
+    parquet_handler.deserialize(data)
 
 
 def test_formats_for_obj():
     arr = np.ndarray(3)
 
-    fmt = FormatRegistry.for_obj(arr)[0]
+    [fmt] = FormatRegistry.for_obj(arr)
 
     assert 'npz' in fmt.handled_extensions
     assert FormatRegistry.for_ext('npy')[0] is fmt
@@ -46,17 +49,17 @@ def test_formats_for_ext():
 
 
 def test_formats_for_meta():
-    bytes_fmt = FormatRegistry.for_meta({'target': 'bytes'})[0]
-    json_fmt = FormatRegistry.for_meta({'target': 'json'})[0]
+    [bytes_fmt] = FormatRegistry.for_meta({'target': 'bytes'})
+    [json_fmt] = FormatRegistry.for_meta({'target': 'json'})
 
     some_bytes = b'["phlipper", "piglet"]'
     assert bytes_fmt.serialize(some_bytes)[0] == some_bytes
     assert json_fmt.deserialize(some_bytes) == ['phlipper', 'piglet']
 
 
 def test_formats_for_format():
-    bytes_fmt = FormatRegistry.for_format('bytes')[0]
-    json_fmt = FormatRegistry.for_format('json')[0]
+    [bytes_fmt] = FormatRegistry.for_format('bytes')
+    [json_fmt] = FormatRegistry.for_format('json')
 
     some_bytes = b'["phlipper", "piglet"]'
     assert bytes_fmt.serialize(some_bytes)[0] == some_bytes
@@ -91,7 +94,7 @@ def test_formats_serdes():
 
 
 def test_formats_csv_read():
-    csv_file = pathlib.Path(__file__).parent / 'data' / 'csv.csv'
+    csv_file = data_dir / 'csv.csv'
 
     meta = {'format': {'name': 'csv'}}
     expected_bytes = b'a,b,c,d\n1,2,3,4\n5,6,7,8\n'
@@ -146,20 +149,23 @@ def test_formats_csv_roundtrip():
     assert df1.equals(df2)
 
 
-def test_formats_search_fail_notfound():
-    # a search that finds nothing should raise with an explanation.
-    class Foo:
-        pass
+def test_formats_anndata():
+    meta = {'format': {'name': 'h5ad'}}
+    ad_file = data_dir / 'test.h5ad'
+    ad = FormatRegistry.deserialize(ad_file.read_bytes(), meta)
+
+    assert isinstance(ad, AnnData)
 
-    bad_kwargs = [
-        dict(obj_type=Foo, meta=None, ext=None),
-        dict(obj_type=None, meta={}, ext=None),
-        dict(obj_type=None, meta=None, ext='.fizz'),
-    ]
 
-    for args in bad_kwargs:
-        with pytest.raises(QuiltException):
-            FormatRegistry.search(**args)
+@pytest.mark.parametrize('args', [
+    dict(obj_type=type('Foo', (), {}), meta=None, ext=None),
+    dict(obj_type=None, meta={}, ext=None),
+    dict(obj_type=None, meta=None, ext='.fizz'),
+])
+def test_formats_search_fail_notfound(args):
+    """a search that finds nothing should raise with an explanation."""
+    with pytest.raises(QuiltException):
+        FormatRegistry.search(**args)
 
 
 def test_formats_search_order():

docs/CHANGELOG.md

@@ -7,6 +7,7 @@
 ## Catalog, Lambdas
 !-->
 ## Python API
+* [Added] Support [AnnData](https://anndata.readthedocs.io/en/latest/) format ([#2974](https://github.com/quiltdata/quilt/pull/2974))
 * [Fixed] Fix check to determine if a file is a tempfile in Windows with Python 3.8+ ([#2900](https://github.com/quiltdata/quilt/pull/2900))
 * [Changed] Disable upload optimization for objects with SSE-KMS ([#2790](https://github.com/quiltdata/quilt/pull/2790))