Merge pull request #226 from astronomy-commons/sean/margin-filter

Refactor `filter_from_pixel_list` from `Catalog` to `HealpixDataset`

src/hipscat/catalog/catalog.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import dataclasses
 from typing import Any, Dict, List, Tuple, Union
 
 import healpy as hp
@@ -121,19 +120,6 @@ def filter_by_polygon(self, vertices: List[SphericalCoordinates] | List[Cartesia
         validate_polygon(vertices)
         return self.filter_from_pixel_list(filter_pixels_by_polygon(self.pixel_tree, vertices))
 
-    def filter_from_pixel_list(self, pixels: List[HealpixPixel]) -> Catalog:
-        """Filter the pixels in the catalog to only include the requested pixels.
-
-        Args:
-            pixels (List[HealpixPixels]): the pixels to include
-
-        Returns:
-            A new catalog with only those pixels. Note that we reset the total_rows
-            to None, instead of performing a scan over the new pixel sizes.
-        """
-        filtered_catalog_info = dataclasses.replace(self.catalog_info, total_rows=None)
-        return Catalog(filtered_catalog_info, pixels)
-
     def generate_negative_tree_pixels(self) -> List[HealpixPixel]:
         """Get the leaf nodes at each healpix order that have zero catalog data.
 

src/hipscat/catalog/healpix_dataset/healpix_dataset.py

@@ -1,7 +1,8 @@
+import dataclasses
 from typing import Any, Dict, List, Tuple, Union
 
 import pandas as pd
-from typing_extensions import TypeAlias
+from typing_extensions import Self, TypeAlias
 
 from hipscat.catalog.dataset import BaseCatalogInfo, Dataset
 from hipscat.catalog.partition_info import PartitionInfo
@@ -102,3 +103,16 @@ def _check_files_exist(cls, catalog_base_dir: FilePointer, storage_options: dict
             raise FileNotFoundError(
                 f"_metadata or partition info file is required in catalog directory {catalog_base_dir}"
             )
+
+    def filter_from_pixel_list(self, pixels: List[HealpixPixel]) -> Self:
+        """Filter the pixels in the catalog to only include the requested pixels.
+
+        Args:
+            pixels (List[HealpixPixels]): the pixels to include
+
+        Returns:
+            A new catalog with only those pixels. Note that we reset the total_rows
+            to None, instead of performing a scan over the new pixel sizes.
+        """
+        filtered_catalog_info = dataclasses.replace(self.catalog_info, total_rows=None)
+        return self.__class__(filtered_catalog_info, pixels)