Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove typing imports for List, Tuple, Union #544

Merged
merged 1 commit into from
Jan 24, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 10 additions & 10 deletions src/lsdb/catalog/catalog.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations

from typing import Callable, List, Tuple, Type
from typing import Callable, Type

import hats as hc
import nested_dask as nd
Expand Down Expand Up @@ -130,7 +130,7 @@ def assign(self, **kwargs) -> Catalog:
def crossmatch(
self,
other: Catalog,
suffixes: Tuple[str, str] | None = None,
suffixes: tuple[str, str] | None = None,
algorithm: (
Type[AbstractCrossmatchAlgorithm] | BuiltInCrossmatchAlgorithm
) = BuiltInCrossmatchAlgorithm.KD_TREE,
Expand Down Expand Up @@ -298,7 +298,7 @@ def cone_search(self, ra: float, dec: float, radius_arcsec: float, fine: bool =
"""
return self.search(ConeSearch(ra, dec, radius_arcsec, fine))

def box_search(self, ra: Tuple[float, float], dec: Tuple[float, float], fine: bool = True) -> Catalog:
def box_search(self, ra: tuple[float, float], dec: tuple[float, float], fine: bool = True) -> Catalog:
"""Performs filtering according to right ascension and declination ranges. The right ascension
edges follow great arc circles and the declination edges follow small arc circles.

Expand Down Expand Up @@ -363,7 +363,7 @@ def order_search(self, min_order: int = 0, max_order: int | None = None) -> Cata
"""
return self.search(OrderSearch(min_order, max_order))

def pixel_search(self, pixels: List[Tuple[int, int]]) -> Catalog:
def pixel_search(self, pixels: list[tuple[int, int]]) -> Catalog:
"""Finds all catalog pixels that overlap with the requested pixel set.

Args:
Expand Down Expand Up @@ -407,12 +407,12 @@ def merge(
self,
other: Catalog,
how: str = "inner",
on: str | List | None = None,
left_on: str | List | None = None,
right_on: str | List | None = None,
on: str | list | None = None,
left_on: str | list | None = None,
right_on: str | list | None = None,
left_index: bool = False,
right_index: bool = False,
suffixes: Tuple[str, str] | None = None,
suffixes: tuple[str, str] | None = None,
) -> nd.NestedFrame:
"""Performs the merge of two catalog Dataframes

Expand Down Expand Up @@ -461,7 +461,7 @@ def merge_asof(
self,
other: Catalog,
direction: str = "backward",
suffixes: Tuple[str, str] | None = None,
suffixes: tuple[str, str] | None = None,
output_catalog_name: str | None = None,
):
"""Uses the pandas `merge_asof` function to merge two catalogs on their indices by distance of keys
Expand Down Expand Up @@ -512,7 +512,7 @@ def join(
left_on: str | None = None,
right_on: str | None = None,
through: AssociationCatalog | None = None,
suffixes: Tuple[str, str] | None = None,
suffixes: tuple[str, str] | None = None,
output_catalog_name: str | None = None,
) -> Catalog:
"""Perform a spatial join to another catalog
Expand Down
4 changes: 1 addition & 3 deletions src/lsdb/catalog/dataset/dataset.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
from typing import List

import hats as hc
import nested_dask as nd
import nested_pandas as npd
Expand Down Expand Up @@ -46,7 +44,7 @@ def compute(self) -> npd.NestedFrame:
"""Compute dask distributed dataframe to pandas dataframe"""
return self._ddf.compute()

def to_delayed(self, optimize_graph: bool = True) -> List[Delayed]:
def to_delayed(self, optimize_graph: bool = True) -> list[Delayed]:
"""Get a list of Dask Delayed objects for each partition in the dataset

Used for more advanced custom operations, but to use again with LSDB, the delayed objects
Expand Down
18 changes: 9 additions & 9 deletions src/lsdb/catalog/dataset/healpix_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

import warnings
from pathlib import Path
from typing import Any, Callable, Dict, Iterable, List, Tuple, Type, cast
from typing import Any, Callable, Iterable, Type, cast

import astropy
import dask
Expand Down Expand Up @@ -118,15 +118,15 @@ def _create_modified_hc_structure(self, **kwargs) -> HCHealpixDataset:
moc=self.hc_structure.moc,
)

def get_healpix_pixels(self) -> List[HealpixPixel]:
def get_healpix_pixels(self) -> list[HealpixPixel]:
"""Get all HEALPix pixels that are contained in the catalog

Returns:
List of all Healpix pixels in the catalog
"""
return self.hc_structure.get_healpix_pixels()

def get_ordered_healpix_pixels(self) -> List[HealpixPixel]:
def get_ordered_healpix_pixels(self) -> list[HealpixPixel]:
"""Get all HEALPix pixels that are contained in the catalog,
ordered by breadth-first nested ordering.

Expand Down Expand Up @@ -189,7 +189,7 @@ def _perform_search(
self,
metadata: hc.catalog.Catalog | hc.catalog.MarginCatalog,
search: AbstractSearch,
) -> Tuple[DaskDFPixelMap, nd.NestedFrame]:
) -> tuple[DaskDFPixelMap, nd.NestedFrame]:
"""Performs a search on the catalog from a list of pixels to search in

Args:
Expand Down Expand Up @@ -237,7 +237,7 @@ def map_partitions(
self,
func: Callable[..., npd.NestedFrame],
*args,
meta: pd.DataFrame | pd.Series | Dict | Iterable | Tuple | None = None,
meta: pd.DataFrame | pd.Series | dict | Iterable | tuple | None = None,
include_pixel: bool = False,
**kwargs,
) -> Self | dd.Series:
Expand Down Expand Up @@ -329,7 +329,7 @@ def prune_empty_partitions(self, persist: bool = False) -> Self:
filtered_hc_structure = self.hc_structure.filter_from_pixel_list(non_empty_pixels)
return self.__class__(search_ddf, ddf_partition_map, filtered_hc_structure)

def _get_non_empty_partitions(self) -> Tuple[List[HealpixPixel], np.ndarray]:
def _get_non_empty_partitions(self) -> tuple[list[HealpixPixel], np.ndarray]:
"""Determines which pixels and partitions of a catalog are not empty

Returns:
Expand All @@ -356,7 +356,7 @@ def skymap_data(
order: int | None = None,
default_value: Any = 0.0,
**kwargs,
) -> Dict[HealpixPixel, Delayed]:
) -> dict[HealpixPixel, Delayed]:
"""Perform a function on each partition of the catalog, returning a dict of values for each pixel.

Args:
Expand Down Expand Up @@ -438,7 +438,7 @@ def skymap(
order: int | None = None,
default_value: Any = 0,
projection="MOL",
plotting_args: Dict | None = None,
plotting_args: dict | None = None,
**kwargs,
) -> tuple[Figure, WCSAxes]:
"""Plot a skymap of an aggregate function applied over each partition
Expand Down Expand Up @@ -722,7 +722,7 @@ def plot_points(
color_col: str | None = None,
projection: str = "MOL",
title: str | None = None,
fov: Quantity | Tuple[Quantity, Quantity] | None = None,
fov: Quantity | tuple[Quantity, Quantity] | None = None,
center: SkyCoord | None = None,
wcs: astropy.wcs.WCS | None = None,
frame_class: Type[BaseFrame] | None = None,
Expand Down
6 changes: 3 additions & 3 deletions src/lsdb/core/crossmatch/abstract_crossmatch_algorithm.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from __future__ import annotations

from abc import ABC
from typing import TYPE_CHECKING, Tuple
from typing import TYPE_CHECKING

import nested_pandas as npd
import numpy as np
Expand Down Expand Up @@ -66,7 +66,7 @@ def __init__(
left_catalog_info: TableProperties,
right_catalog_info: TableProperties,
right_margin_catalog_info: TableProperties | None,
suffixes: Tuple[str, str],
suffixes: tuple[str, str],
):
"""Initializes a crossmatch algorithm

Expand Down Expand Up @@ -107,7 +107,7 @@ def crossmatch(self, **kwargs) -> npd.NestedFrame:
)
return self._create_crossmatch_df(l_inds, r_inds, extra_cols)

def perform_crossmatch(self) -> Tuple[np.ndarray, np.ndarray, pd.DataFrame]:
def perform_crossmatch(self) -> tuple[np.ndarray, np.ndarray, pd.DataFrame]:
"""Performs a crossmatch to get the indices of the matching rows and any extra columns

Any additional keyword arguments needed can be added to this method in the subclass, and the user
Expand Down
4 changes: 2 additions & 2 deletions src/lsdb/core/crossmatch/bounded_kdtree_match.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations

from typing import TYPE_CHECKING, Tuple
from typing import TYPE_CHECKING

import numpy as np
import pandas as pd
Expand Down Expand Up @@ -36,7 +36,7 @@ def perform_crossmatch(
n_neighbors: int = 1,
radius_arcsec: float = 1,
min_radius_arcsec: float = 0,
) -> Tuple[np.ndarray, np.ndarray, pd.DataFrame]:
) -> tuple[np.ndarray, np.ndarray, pd.DataFrame]:
"""Perform a cross-match between the data from two HEALPix pixels

Finds the n closest neighbors in the right catalog for each point in the left catalog that
Expand Down
6 changes: 3 additions & 3 deletions src/lsdb/core/crossmatch/kdtree_match.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations

from typing import TYPE_CHECKING, Tuple
from typing import TYPE_CHECKING

import numpy as np
import numpy.typing as npt
Expand Down Expand Up @@ -44,7 +44,7 @@ def perform_crossmatch(
self,
n_neighbors: int = 1,
radius_arcsec: float = 1,
) -> Tuple[np.ndarray, np.ndarray, pd.DataFrame]:
) -> tuple[np.ndarray, np.ndarray, pd.DataFrame]:
"""Perform a cross-match between the data from two HEALPix pixels

Finds the n closest neighbors in the right catalog for each point in the left catalog that
Expand Down Expand Up @@ -72,7 +72,7 @@ def perform_crossmatch(
)
return left_idx, right_idx, extra_columns

def _get_point_coordinates(self) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
def _get_point_coordinates(self) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
left_xyz = _lon_lat_to_xyz(
lon=self.left[self.left_catalog_info.ra_column].to_numpy(),
lat=self.left[self.left_catalog_info.dec_column].to_numpy(),
Expand Down
5 changes: 2 additions & 3 deletions src/lsdb/core/crossmatch/kdtree_utils.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import math
from typing import Tuple

import numpy as np
import numpy.typing as npt
Expand All @@ -12,7 +11,7 @@ def _find_crossmatch_indices(
n_neighbors: int,
max_distance: float,
min_distance: float = 0,
) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.int64], npt.NDArray[np.int64]]:
) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.int64], npt.NDArray[np.int64]]:
# If right catalog is empty, tree.query will raise an exception
# Left catalog cannot be empty, as it checked in perform_crossmatch()
if len(right_xyz) == 0:
Expand Down Expand Up @@ -52,7 +51,7 @@ def _query_min_max_neighbors(
n_neighbors: int,
min_distance: float,
max_distance: float,
) -> Tuple[np.ndarray, np.ndarray]:
) -> tuple[np.ndarray, np.ndarray]:
"""Finds `n_neighbors` within a distance range for all points in a pair of partitions"""
left_tree = KDTree(left_xyz, compact_nodes=True, balanced_tree=True, copy_data=False)

Expand Down
4 changes: 2 additions & 2 deletions src/lsdb/core/plotting/plot_points.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations

from typing import Tuple, Type
from typing import Type

import astropy
import matplotlib.pyplot as plt
Expand All @@ -22,7 +22,7 @@ def plot_points(
color_col: str | None = None,
projection: str = "MOL",
title: str = "",
fov: Quantity | Tuple[Quantity, Quantity] | None = None,
fov: Quantity | tuple[Quantity, Quantity] | None = None,
center: SkyCoord | None = None,
wcs: astropy.wcs.WCS | None = None,
frame_class: Type[BaseFrame] | None = None,
Expand Down
4 changes: 2 additions & 2 deletions src/lsdb/core/plotting/skymap.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations

from typing import Any, Callable, Dict
from typing import Any, Callable

import hats.pixel_math.healpix_shim as hp
import nested_pandas as npd
Expand Down Expand Up @@ -40,7 +40,7 @@ def apply_func(df):


def compute_skymap(
pixel_map: Dict[HealpixPixel, Any], order: int | None = None, default_value: Any = 0.0
pixel_map: dict[HealpixPixel, Any], order: int | None = None, default_value: Any = 0.0
) -> np.ndarray:
"""Returns a histogram map of healpix_pixels to values.

Expand Down
4 changes: 2 additions & 2 deletions src/lsdb/core/search/abstract_search.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from __future__ import annotations

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Tuple, Type
from typing import TYPE_CHECKING, Type

import astropy
import nested_pandas as npd
Expand Down Expand Up @@ -45,7 +45,7 @@ def plot(
self,
projection: str = "MOL",
title: str = "",
fov: Quantity | Tuple[Quantity, Quantity] | None = None,
fov: Quantity | tuple[Quantity, Quantity] | None = None,
center: SkyCoord | None = None,
wcs: astropy.wcs.WCS | None = None,
frame_class: Type[BaseFrame] | None = None,
Expand Down
4 changes: 2 additions & 2 deletions src/lsdb/core/search/pixel_search.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations

from typing import TYPE_CHECKING, List, Tuple
from typing import TYPE_CHECKING

import nested_pandas as npd
from hats.pixel_math import HealpixPixel
Expand All @@ -18,7 +18,7 @@ class PixelSearch(AbstractSearch):
Does not filter points inside those partitions.
"""

def __init__(self, pixels: List[Tuple[int, int]]):
def __init__(self, pixels: list[tuple[int, int]]):
super().__init__(fine=False)
self.pixels = [HealpixPixel(o, p) for o, p in set(pixels)]

Expand Down
6 changes: 3 additions & 3 deletions src/lsdb/dask/crossmatch_catalog_data.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from __future__ import annotations

import warnings
from typing import TYPE_CHECKING, Tuple, Type
from typing import TYPE_CHECKING, Type

import dask
import nested_dask as nd
Expand Down Expand Up @@ -76,12 +76,12 @@ def perform_crossmatch(
def crossmatch_catalog_data(
left: Catalog,
right: Catalog,
suffixes: Tuple[str, str],
suffixes: tuple[str, str],
algorithm: (
Type[AbstractCrossmatchAlgorithm] | BuiltInCrossmatchAlgorithm
) = BuiltInCrossmatchAlgorithm.KD_TREE,
**kwargs,
) -> Tuple[nd.NestedFrame, DaskDFPixelMap, PixelAlignment]:
) -> tuple[nd.NestedFrame, DaskDFPixelMap, PixelAlignment]:
"""Cross-matches the data from two catalogs

Args:
Expand Down
4 changes: 1 addition & 3 deletions src/lsdb/dask/divisions.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,12 @@
from __future__ import annotations

from typing import List, Tuple

import numpy as np
from hats.pixel_math import HealpixPixel
from hats.pixel_math.healpix_pixel_function import get_pixel_argsort
from hats.pixel_math.spatial_index import healpix_to_spatial_index


def get_pixels_divisions(healpix_pixels: List[HealpixPixel]) -> Tuple[int, ...] | None:
def get_pixels_divisions(healpix_pixels: list[HealpixPixel]) -> tuple[int, ...] | None:
"""Calculates the Dask Dataframe divisions for a list of HEALPix pixels.

Divisions include the minimum value of every HEALPix pixel spatial_index
Expand Down
Loading
Loading