From 87655bc59894114022b5845013917de552ea59b0 Mon Sep 17 00:00:00 2001
From: colganwi <colgan.william@gmail.com>
Date: Mon, 2 Sep 2024 10:03:43 -0400
Subject: [PATCH] typing import patch

---
 src/treedata/_core/treedata.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/treedata/_core/treedata.py b/src/treedata/_core/treedata.py
index 4e18767..d9ff34d 100755
--- a/src/treedata/_core/treedata.py
+++ b/src/treedata/_core/treedata.py
@@ -16,7 +16,7 @@
 import numpy as np
 import pandas as pd
 import zarr
-from anndata._core.index import Index, Index1D, _subset
+from anndata._core.index import _subset
 from anndata._io import write_h5ad, write_zarr
 from scipy import sparse
 
@@ -29,6 +29,9 @@
 if TYPE_CHECKING:
     from os import PathLike
 
+    Index1D = slice | int | str | np.int64 | np.ndarray
+    Index = Index1D | tuple[Index1D, Index1D] | sparse.spmatrix | sparse.sparray
+
 
 class TreeData(ad.AnnData):
     """AnnData with trees.