From 0330f846b2385f6f87edda3f3d03edcccce34a5c Mon Sep 17 00:00:00 2001
From: Sandro Campos <scampos@andrew.cmu.edu>
Date: Fri, 27 Sep 2024 16:00:28 -0400
Subject: [PATCH] Add import for annotations

---
 benchmarks/benchmarks.py | 4 +++-
 src/corrgi/alignment.py  | 3 +++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/benchmarks/benchmarks.py b/benchmarks/benchmarks.py
index 0f652d2..d4a1f82 100644
--- a/benchmarks/benchmarks.py
+++ b/benchmarks/benchmarks.py
@@ -51,7 +51,9 @@ def time_pcf_natural_estimator(self):
 
     def time_pccf_davis_peebles_estimator(self):
         """Times the Davis-Peebles estimator for a projected cross-correlation"""
-        estimator = DavisPeeblesEstimator(ProjectedCorrelation(params=generate_pcf_params(), use_weights=True))
+        estimator = DavisPeeblesEstimator(
+            ProjectedCorrelation(params=generate_pcf_params(), use_weights=True)
+        )
         estimator.compute_crosscorrelation_counts(
             GALS_WEIGHT_DIR,
             GALS1_WEIGHT_DIR,
diff --git a/src/corrgi/alignment.py b/src/corrgi/alignment.py
index 4496e8d..36277c0 100644
--- a/src/corrgi/alignment.py
+++ b/src/corrgi/alignment.py
@@ -1,3 +1,4 @@
+from __future__ import annotations
 import itertools
 import pickle
 
@@ -57,6 +58,7 @@ def crosscorrelation_alignment(catalog_left: Catalog, catalog_right: Catalog) ->
     result_mapping = pd.DataFrame(full_product, columns=column_names)
     return PixelAlignment(catalog_left.pixel_tree, result_mapping, PixelAlignmentType.OUTER)
 
+
 def read_alignment(path: str) -> dict | list:
     """Reads the alignment pixels from a file.
 
@@ -71,6 +73,7 @@ def read_alignment(path: str) -> dict | list:
     with open(path, "rb") as alignment_file:
         return pickle.load(alignment_file)
 
+
 def write_alignment(path: str, pixels: dict | list):
     """Writes the alignment pixels to a file.