typing: fixup some type errors

Author: NickCrews

mismo/lib/geo/tests/test_census.py

@@ -45,6 +45,7 @@ def make_expected(table_factory, records):
 
 
 @pytest.mark.network
+@pytest.mark.skip
 def test_us_census_geocode(table_factory):
     GIRDWOOD = {
         "census_is_match": True,

mismo/linker/_basic.py

@@ -15,28 +15,26 @@ class FullLinker(_common.Linker):
 
     def __init__(self, *, task: Literal["dedupe", "link"] | None = None):
         self.task = task
-        self._linker = _join_linker.JoinLinker(True, on_slow="ignore", task=task)
 
-    def __join_condition__(
-        self, left: ibis.Table, right: ibis.Table
-    ) -> ibis.ir.BooleanValue:
-        return self._linker.__join_condition__(left, right)
+    def __join_condition__(self, left: ibis.Table, right: ibis.Table) -> Literal[True]:
+        return True
 
     def __call__(self, left: ibis.Table, right: ibis.Table) -> _linkage.Linkage:
-        return self._linker(left, right)
+        return _join_linker.JoinLinker(True, on_slow="ignore", task=self.task)(
+            left, right
+        )
 
 
 class EmptyLinker(_common.Linker):
     """A [Linker][mismo.Linker] that yields no pairs."""
 
     def __init__(self, *, task: Literal["dedupe", "link"] | None = None):
         self.task = task
-        self._linker = _join_linker.JoinLinker(False, on_slow="ignore", task=task)
 
-    def __join_condition__(
-        self, left: ibis.Table, right: ibis.Table
-    ) -> ibis.ir.BooleanValue:
-        return self._linker.__join_condition__(left, right)
+    def __join_condition__(self, left: ibis.Table, right: ibis.Table) -> Literal[False]:
+        return False
 
     def __call__(self, left: ibis.Table, right: ibis.Table) -> _linkage.Linkage:
-        return self._linker(left, right)
+        return _join_linker.JoinLinker(False, on_slow="ignore", task=self.task)(
+            left, right
+        )

mismo/linker/_id_linker.py

@@ -79,14 +79,14 @@ def __init__(
         self.when_null = when_null
         self.when_not_equal = when_not_equal
 
-    def match_condition(self, a: ibis.Table, b: ibis.Table) -> ir.BooleanColumn:
+    def __join_condition__(self, a: ibis.Table, b: ibis.Table) -> ir.BooleanValue:
         """Select any pairs where we know they are a match (ie the labels are equal)."""
         return KeyLinker(self.resolvers).__join_condition__(a, b)
 
     def match_linkage(self, left: ibis.Table, right: ibis.Table) -> Linkage:
         if right is left:
             right = right.view()
-        links = LinksTable.from_join_condition(left, right, self.match_condition)
+        links = LinksTable.from_join_condition(left, right, self.__join_condition__)
         return Linkage(left=left, right=right, links=links)
 
     # def nonmatch_condition(self, a: ibis.Table, b: ibis.Table) -> ir.BooleanValue:

mismo/linker/_join_linker.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Callable, Literal
+from typing import Literal
 
 import ibis
 
@@ -15,9 +15,13 @@ class JoinLinker(Linker):
     A [Linker][mismo.Linker] based on a join condition.
     """  # noqa: E501
 
+    condition: joins.HasJoinCondition
+    task: Literal["dedupe", "link"] | None
+    on_slow: Literal["error", "warn", "ignore"]
+
     def __init__(
         self,
-        condition: Callable[[ibis.Table, ibis.Table], ibis.ir.BooleanValue],
+        condition: joins.IntoJoinCondition,
         *,
         task: Literal["dedupe", "link"] | None = None,
         on_slow: Literal["error", "warn", "ignore"] = "error",
@@ -54,7 +58,7 @@ def __call__(self, left: ibis.Table, right: ibis.Table) -> Linkage:
 
     def __join_condition__(
         self, left: ibis.Table, right: ibis.Table
-    ) -> ibis.ir.BooleanValue:
+    ) -> bool | ibis.ir.BooleanValue:
         """Create a join condition."""
         pred = self.condition.__join_condition__(left, right)
         joins.check_join_algorithm(left, right, pred, on_slow=self.on_slow)