refactor(typing): Include all suffixes in Extension

dangotbanned · dangotbanned · commit 755ab4f560af · 2025-01-30T19:45:00.000Z
Also simplifies and removes outdated `Extension`-related tooling
diff --git a/altair/datasets/_reader.py b/altair/datasets/_reader.py
@@ -41,7 +41,6 @@
     module_not_found,
 )
 from altair.datasets._readimpl import IntoFrameT, is_available
-from altair.datasets._typing import EXTENSION_SUFFIXES
 
 if TYPE_CHECKING:
     import sys
@@ -443,8 +442,12 @@ def _into_constraints(
     elif suffix.startswith("."):
         m = {"dataset_name": name, "suffix": suffix}
     else:
+        from typing import get_args
+
+        from altair.datasets._typing import Extension
+
         msg = (
-            f"Expected 'suffix' to be one of {EXTENSION_SUFFIXES!r},\n"
+            f"Expected 'suffix' to be one of {get_args(Extension)!r},\n"
             f"but got: {suffix!r}"
         )
         raise TypeError(msg)
diff --git a/altair/datasets/_typing.py b/altair/datasets/_typing.py
@@ -4,25 +4,20 @@
 from __future__ import annotations
 
 import sys
-from typing import Any, Literal
+from typing import Literal
 
 if sys.version_info >= (3, 14):
     from typing import TypedDict
 else:
     from typing_extensions import TypedDict
 
-if sys.version_info >= (3, 13):
-    from typing import TypeIs
-else:
-    from typing_extensions import TypeIs
-
 if sys.version_info >= (3, 10):
     from typing import TypeAlias
 else:
     from typing_extensions import TypeAlias
 
 
-__all__ = ["EXTENSION_SUFFIXES", "Dataset", "Extension", "Metadata", "is_ext_read"]
+__all__ = ["Dataset", "Extension", "Metadata"]
 
 Dataset: TypeAlias = Literal[
     "7zip",
@@ -98,18 +93,7 @@
     "world-110m",
     "zipcodes",
 ]
-Extension: TypeAlias = Literal[".arrow", ".csv", ".json", ".parquet", ".tsv"]
-EXTENSION_SUFFIXES: tuple[
-    Literal[".arrow"],
-    Literal[".csv"],
-    Literal[".json"],
-    Literal[".parquet"],
-    Literal[".tsv"],
-] = (".arrow", ".csv", ".json", ".parquet", ".tsv")
-
-
-def is_ext_read(suffix: Any) -> TypeIs[Extension]:
-    return suffix in {".arrow", ".csv", ".json", ".parquet", ".tsv"}
+Extension: TypeAlias = Literal[".arrow", ".csv", ".json", ".parquet", ".png", ".tsv"]
 
 
 class Metadata(TypedDict, total=False):
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
@@ -16,7 +16,7 @@
 
 from altair.datasets import Loader
 from altair.datasets._exceptions import AltairDatasetsError
-from altair.datasets._typing import Dataset, Metadata, is_ext_read
+from altair.datasets._typing import Dataset, Metadata
 from tests import no_xdist, skip_requires_pyarrow
 from tools import fs
 
@@ -441,10 +441,7 @@ def test_reader_cache_exhaustive(
 
     # NOTE: Approximating all datasets downloaded
     assert len(cached_paths) >= 70
-    assert all(
-        bool(fp.exists() and is_ext_read(fp.suffix) and fp.stat().st_size)
-        for fp in load.cache
-    )
+    assert all(bool(fp.exists() and fp.stat().st_size) for fp in load.cache)
     # NOTE: Confirm this is a no-op
     load.cache.download_all()
     assert len(cached_paths) == len(tuple(load.cache))
diff --git a/tools/datasets/__init__.py b/tools/datasets/__init__.py
@@ -190,13 +190,6 @@ def generate_typing(self, dpkg: datapackage.DataPackage) -> None:
         indent = " " * 4
         NAME = "Dataset"
         EXT = "Extension"
-        EXT_TYPES = dpkg.extensions()
-        EXTENSION_SUFFIXES = "EXTENSION_SUFFIXES"
-        EXTENSION_TYPE_TP = (
-            f"tuple[{', '.join(f'Literal[{el!r}]' for el in EXT_TYPES)}]"
-        )
-        EXTENSION_GUARD = "is_ext_read"
-
         FIELD = "FlFieldStr"
         FIELD_TYPES = (
             "integer",
@@ -215,17 +208,13 @@ def generate_typing(self, dpkg: datapackage.DataPackage) -> None:
             f"{HEADER_COMMENT}",
             "from __future__ import annotations\n",
             "import sys",
-            "from typing import Any, Literal, TYPE_CHECKING",
+            "from typing import Literal, TYPE_CHECKING",
             utils.import_typing_extensions((3, 14), "TypedDict"),
-            utils.import_typing_extensions((3, 13), "TypeIs"),
             utils.import_typing_extensions((3, 10), "TypeAlias"),
             "\n",
-            f"__all__ = {[NAME, EXT, dpkg._NAME_TYPED_DICT, EXTENSION_GUARD, EXTENSION_SUFFIXES]}\n",
+            f"__all__ = {[NAME, EXT, dpkg._NAME_TYPED_DICT]}\n",
             utils.spell_literal_alias(NAME, dpkg.dataset_names()),
-            utils.spell_literal_alias(EXT, EXT_TYPES),
-            f"{EXTENSION_SUFFIXES}: {EXTENSION_TYPE_TP} = {EXT_TYPES!r}",
-            f"def {EXTENSION_GUARD}(suffix: Any) -> TypeIs[{EXT}]:\n"
-            f"{indent}return suffix in set({EXT_TYPES!r})\n",
+            utils.spell_literal_alias(EXT, dpkg.extensions()),
             dpkg.typed_dict(),
             utils.spell_literal_alias(FIELD, FIELD_TYPES),
             '"""\n'
diff --git a/tools/datasets/datapackage.py b/tools/datasets/datapackage.py
@@ -103,8 +103,7 @@ def dataset_names(self) -> Iterable[str]:
 
     def extensions(self) -> tuple[str, ...]:
         return tuple(
-            self.core.filter(is_image=False)
-            .select(col("suffix").unique().sort())
+            self.core.select(col("suffix").unique().sort())
             .collect()
             .to_series()
             .to_list()
