generics working

README.md

@@ -74,12 +74,10 @@ We are unable to support the functions that work with geometries.
 | Function | Description | Supported|
 | --: | --- | ---|
 | `latlng_to_cell` | Convert latitude/longitude coordinate to cell ID | ✅|
-| `latlng_to_cell_string` | Convert latitude/longitude coordinate to cell ID (returns VARCHAR) | ✅ |
 | `cell_to_lat` | Convert cell ID to latitude | ✅ |
 | `cell_to_lng` | Convert cell ID to longitude | ✅ |
 | `cell_to_latlng` | Convert cell ID to latitude/longitude | ✅ |
 | `get_resolution` | Get resolution number of cell ID | ✅ |
-| `get_base_cell_number` | Get base cell number of cell ID | 🚧|
 | `str_to_int` | Convert VARCHAR cell ID to UBIGINT | ✅ |
 | `int_to_str` | Convert BIGINT or UBIGINT cell ID to VARCHAR | ✅ |
 | `is_valid_cell` | True if this is a valid cell ID | ✅ |
@@ -110,7 +108,6 @@ We are unable to support the functions that work with geometries.
 | `get_directed_edge_destination` | Convert a directed edge ID to destination cell ID | ✅|
 | `cells_to_directed_edge` | Convert an origin/destination pair to directed edge ID | ✅ |
 | `are_neighbor_cells` | True if the two cell IDs are directly adjacent | ✅ |
-| `directed_edge_to_boundary_wkt` | Convert directed edge ID to linestring WKT | ✅ |
 | `average_hexagon_area` | Get average area of a hexagon cell at resolution |  ✅ |
 | `cell_area` | Get the area of a cell ID |  ✅|
 | `average_hexagon_edge_length` | Average hexagon edge length at resolution |  ✅|
@@ -122,3 +119,4 @@ We are unable to support the functions that work with geometries.
 | `cells_to_multi_polygon_wkt` | Convert a set of cells to multipolygon WKT | 🛑 |
 | `polygon_wkt_to_cells` | Convert polygon WKT to a set of cells | 🛑 |
 | `cell_to_boundary_wkt` | Convert cell ID to cell boundary | 🛑 |
+| `directed_edge_to_boundary_wkt` | Convert directed edge ID to linestring WKT | 🛑 |

polars_h3/__init__.py

@@ -1,68 +1,65 @@
-from .core.traversal import (
-    grid_distance,
-    grid_ring,
-    grid_disk,
-    grid_path_cells,
+from .core.edge import (
+    are_neighbor_cells,
+    cells_to_directed_edge,
+    directed_edge_to_boundary,
+    directed_edge_to_cells,
+    get_directed_edge_destination,
+    get_directed_edge_origin,
+    is_valid_directed_edge,
+    origin_to_directed_edges,
 )
 from .core.indexing import (
-    latlng_to_cell,
-    latlng_to_cell_string,
     cell_to_lat,
-    cell_to_lng,
     cell_to_latlng,
+    cell_to_lng,
     cell_to_local_ij,
+    latlng_to_cell,
     local_ij_to_cell,
 )
 from .core.inspection import (
-    get_resolution,
-    str_to_int,
-    int_to_str,
-    is_valid_cell,
-    is_pentagon,
-    is_res_class_III,
-    get_icosahedron_faces,
-    cell_to_parent,
     cell_to_center_child,
-    cell_to_children_size,
-    cell_to_children,
     cell_to_child_pos,
+    cell_to_children,
+    cell_to_children_size,
+    cell_to_parent,
     child_pos_to_cell,
     compact_cells,
+    get_icosahedron_faces,
+    get_resolution,
+    int_to_str,
+    is_pentagon,
+    is_res_class_III,
+    is_valid_cell,
+    str_to_int,
     uncompact_cells,
 )
-from .core.vertexes import (
-    cell_to_vertex,
-    cell_to_vertexes,
-    vertex_to_latlng,
-    is_valid_vertex,
-)
-from .core.edge import (
-    are_neighbor_cells,
-    cells_to_directed_edge,
-    is_valid_directed_edge,
-    get_directed_edge_origin,
-    get_directed_edge_destination,
-    directed_edge_to_cells,
-    origin_to_directed_edges,
-    directed_edge_to_boundary,
-)
 from .core.metrics import (
-    great_circle_distance,
     average_hexagon_area,
+    average_hexagon_edge_length,
     cell_area,
     edge_length,
-    average_hexagon_edge_length,
     get_num_cells,
+    great_circle_distance,
+)
+from .core.traversal import (
+    grid_disk,
+    grid_distance,
+    grid_path_cells,
+    grid_ring,
+)
+from .core.vertexes import (
+    cell_to_vertex,
+    cell_to_vertexes,
+    is_valid_vertex,
+    vertex_to_latlng,
 )
-
 
 __all__ = [
     "grid_distance",
     "grid_ring",
     "grid_disk",
     "grid_path_cells",
     "latlng_to_cell",
-    "latlng_to_cell_string",
     "cell_to_lat",
     "cell_to_lng",
     "cell_to_latlng",

polars_h3/core/_types.py

@@ -0,0 +1,40 @@
+from typing import Literal, TypeVar, Union
+
+from typing_extensions import override
+
+HexResolution = Union[
+    Literal[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], int
+]
+
+_T = TypeVar("_T")
+
+
+# Sentinel class used until PEP 0661 is accepted
+class NotGiven:
+    """
+    A sentinel singleton class used to distinguish omitted keyword arguments
+    from those passed in with the value None (which may have different behavior).
+
+    For example:
+
+    ```py
+    def get(timeout: Union[int, NotGiven, None] = NotGiven()) -> Response:
+        ...
+
+
+    get(timeout=1)  # 1s timeout
+    get(timeout=None)  # No timeout
+    get()  # Default timeout behavior, which may not be statically known at the method definition.
+    ```
+    """
+
+    def __bool__(self) -> Literal[False]:
+        return False
+
+    @override
+    def __repr__(self) -> str:
+        return "NOT_GIVEN"
+
+
+NotGivenOr = Union[_T, NotGiven]
+NOT_GIVEN = NotGiven()

polars_h3/core/edge.py

@@ -1,12 +1,11 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import polars as pl
 from polars.plugins import register_plugin_function
 
-
 if TYPE_CHECKING:
     from polars_h3.typing import IntoExprColumn
 

polars_h3/core/indexing.py

@@ -15,20 +15,37 @@
 LIB = Path(__file__).parent.parent
 
 
-def latlng_to_cell_string(
-    lat: IntoExprColumn, lng: IntoExprColumn, resolution: HexResolution
+def latlng_to_cell(
+    lat: IntoExprColumn,
+    lng: IntoExprColumn,
+    resolution: HexResolution,
+    return_dtype: type[pl.DataType] = pl.UInt64,
 ) -> pl.Expr:
     """
-    Indexes the location at the specified resolution, providing the index of the cell containing the location. This buckets the geographic point into the H3 grid.
+    Indexes the location at the specified resolution, providing the index of the cell containing the location. The result dtype is determined by `return_dtype`.
     """
     assert_valid_resolution(resolution)
-    return register_plugin_function(
-        args=[lat, lng],
-        plugin_path=LIB,
-        function_name="latlng_to_cell_string",
-        is_elementwise=True,
-        kwargs={"resolution": resolution},
-    )
+
+    if return_dtype == pl.Utf8:
+        expr = register_plugin_function(
+            args=[lat, lng],
+            plugin_path=LIB,
+            function_name="latlng_to_cell_string",
+            is_elementwise=True,
+            kwargs={"resolution": resolution},
+        )
+    else:
+        expr = register_plugin_function(
+            args=[lat, lng],
+            plugin_path=LIB,
+            function_name="latlng_to_cell",
+            is_elementwise=True,
+            kwargs={"resolution": resolution},
+        )
+        if return_dtype != pl.UInt64:
+            expr = expr.cast(return_dtype)
+
+    return expr
 
 
 def cell_to_lat(cell: IntoExprColumn) -> pl.Expr:

polars_h3/core/inspection.py

@@ -1,13 +1,12 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import TYPE_CHECKING, Union
-from pathlib import Path
+from typing import TYPE_CHECKING
 
 import polars as pl
 from polars.plugins import register_plugin_function
-from .utils import _assert_valid_resolution, HexResolution
 
+from .utils import HexResolution, assert_valid_resolution
 
 if TYPE_CHECKING:
     from polars_h3.typing import IntoExprColumn
@@ -101,7 +100,8 @@ def get_icosahedron_faces(expr: IntoExprColumn) -> pl.Expr:
 
 
 def cell_to_parent(
-    cell: IntoExprColumn, resolution: Union[HexResolution, None] = None
+    cell: IntoExprColumn,
+    resolution: HexResolution,
 ) -> pl.Expr:
     """
     Returns the parent cell at the specified resolution. If the input cell has resolution r, then parentRes = r - 1 would give the immediate parent, parentRes = r - 2 would give the grandparent, and so on.

polars_h3/core/traversal.py

@@ -1,12 +1,11 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import polars as pl
 from polars.plugins import register_plugin_function
 
-
 if TYPE_CHECKING:
     from polars_h3.typing import IntoExprColumn
 

polars_h3/core/utils.py

@@ -1,6 +1,4 @@
-from typing import Literal
-
-HexResolution = Literal[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
+from ._types import HexResolution
 
 
 def assert_valid_resolution(resolution: HexResolution) -> None:

polars_h3/core/vertexes.py

@@ -1,12 +1,11 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import polars as pl
 from polars.plugins import register_plugin_function
 
-
 if TYPE_CHECKING:
     from polars_h3.typing import IntoExprColumn