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 = []