|
import enum |
|
import json |
|
from collections.abc import Generator, Iterable, MutableSequence |
|
from typing import Annotated, TypedDict, cast |
|
|
|
import pyogrio |
|
import typer |
|
from rich.console import Console |
|
from rich.table import Table |
|
|
|
_APP = typer.Typer() |
|
|
|
|
|
class DriverDict(TypedDict): |
|
"""Represents a GDAL driver""" |
|
|
|
name: str |
|
long_name: str |
|
read: bool |
|
append: bool |
|
write: bool |
|
supports_vsi: bool |
|
help_topic_url: str | None |
|
extensions: list[str] | None |
|
|
|
|
|
type DriversDict = dict[str, DriverDict] |
|
"""Output of pyogrio.list_drivers_details() |
|
""" |
|
|
|
|
|
class OptionsDict(TypedDict, total=False): |
|
read: bool |
|
append: bool |
|
write: bool |
|
supports_vsi: bool |
|
help_topic_url: str |
|
extensions: Iterable[str] |
|
|
|
|
|
def meets_criteria( |
|
d: DriverDict, |
|
read: bool | None = None, |
|
append: bool | None = None, |
|
write: bool | None = None, |
|
supports_vsi: bool | None = None, |
|
extension: str | None = None, |
|
): |
|
"""Determines if a driver has specified capabilities. |
|
Only capabilities parameters that are not None will be checked. |
|
|
|
Args: |
|
d: A dict representing a GDAL driver. |
|
read Check to see if a driver supports reading. Defaults to None. |
|
append: Check to see if a driver supports appending. Defaults to None. |
|
write: Check to see if a driver supports writing. Defaults to None. |
|
supports_vsi: Set to True to check for VSI support. Defaults to None. |
|
extension: Specify an extension to check for. |
|
Comparison of extensions is case-insensitive. Defaults to None. |
|
|
|
Returns: |
|
True if the criteria are met, false otherwise. |
|
""" |
|
|
|
def is_valid(a: bool | None, b: bool): |
|
"""Returns True if: |
|
|
|
* A is None or False |
|
* A is non-None and b is True |
|
|
|
Args: |
|
a (bool | None): A boolean value or None |
|
b (bool): A boolean value |
|
|
|
Returns: |
|
bool: True or False |
|
""" |
|
if not a: |
|
return True |
|
return b |
|
|
|
# Get bool values for the given parameters |
|
# and those of the corresponding values of in the dict |
|
booleans = [ |
|
(read, d["read"]), |
|
(append, d["append"]), |
|
(write, d["write"]), |
|
(supports_vsi, d["supports_vsi"]), |
|
] |
|
|
|
# Test all of the capabilities for which the |
|
# corresponding parameter was True. |
|
tests = (is_valid(*n) for n in booleans) |
|
|
|
# Return False if any of the tests were False |
|
if not all(tests): |
|
return False |
|
|
|
# If there was an extension provided, return True only |
|
# if the driver supports that extension. |
|
if extension: |
|
driver_extentions = d["extensions"] |
|
|
|
# Return false if the user provided an extension |
|
# but the driver does not have any extensions |
|
if not driver_extentions: |
|
return False |
|
|
|
# Add a leading period if the extension does not already have one. |
|
if not extension.startswith("."): |
|
extension = f".{extension}" |
|
|
|
# Return True if the driver is in the list of drivers, false otherwise. |
|
return extension.casefold() in (d.casefold() for d in driver_extentions) |
|
|
|
# Return True if nothing else has failed or returned a value. |
|
return True |
|
|
|
|
|
def get_flattened_drivers_list( |
|
read: bool | None = None, |
|
append: bool | None = None, |
|
write: bool | None = None, |
|
supports_vsi: bool | None = None, |
|
extension: str | None = None, |
|
): |
|
"""Get a list of supported GDAL drivers |
|
|
|
Args: |
|
read: Supports reading. Defaults to None. |
|
append: Support appending. Defaults to None. |
|
write: Supports writing. Defaults to None. |
|
supports_vsi: Supports VSI. Defaults to None. |
|
extension: Supports the given extension. Defaults to None. |
|
|
|
Returns: |
|
A list of drivers that meet the given criteria |
|
""" |
|
|
|
# pyogrio.list_drivers_details returns a |
|
# dict[str, dict] keyed by the driver's name. |
|
# |
|
# 1. Add the driver name to each child dict. |
|
# 2. Return a list of these child dicts, now with a "name". |
|
output = ({**{"name": k}, **v} for k, v in (pyogrio.list_drivers_details()).items()) |
|
output = cast(Generator[DriverDict], output) |
|
|
|
# Filter out the drivers that do not support the features |
|
# specified by the function parameters. |
|
output = [ |
|
d |
|
for d in output |
|
if meets_criteria(d, read, append, write, supports_vsi, extension) |
|
] |
|
return output |
|
|
|
|
|
# We don't know ahead of time what drivers will be available, so we have to create the enum dynamically. |
|
VectorDriverNames = enum.Enum( |
|
"DriversEnum", |
|
((d["name"], d["name"]) for d in get_flattened_drivers_list() if d["write"]), |
|
) |
|
"""An enumeration of valid vector driver names |
|
""" |
|
|
|
|
|
def has_extension(d: DriverDict, ext: str): |
|
extension_list = d["extensions"] |
|
if not extension_list: |
|
return False |
|
extension_list = [x.casefold() == ext.casefold() for x in extension_list] |
|
return any(extension_list) |
|
|
|
|
|
@_APP.command() |
|
def main( |
|
read: Annotated[bool, typer.Option("--read", "-r")] = False, |
|
write: Annotated[bool, typer.Option("--write", "-w")] = False, |
|
append: Annotated[bool, typer.Option("--append", "-a")] = False, |
|
supports_vsi: Annotated[ |
|
bool, typer.Option("--supports-vsi", "--vsi", "-v") |
|
] = False, |
|
as_json: Annotated[bool, typer.Option()] = False, |
|
extension: Annotated[str | None, typer.Option(case_sensitive=False)] = None, |
|
): |
|
"""Returns all drivers that match the given criteria""" |
|
|
|
with Console() as c: |
|
driver_dicts = get_flattened_drivers_list( |
|
read, write, append, supports_vsi, extension |
|
) |
|
|
|
# Sort by "name" |
|
driver_dicts.sort(key=lambda d: d["name"]) |
|
|
|
if as_json: |
|
c.print_json(json.dumps(driver_dicts)) |
|
else: |
|
|
|
def format_heading(s: str): |
|
parts = s.split("_") |
|
if parts[0] == "help": |
|
return parts[0] |
|
if parts[-1] == "vsi": |
|
return parts[-1] |
|
return " ".join(parts) |
|
|
|
table_headings = (format_heading(k) for k in driver_dicts[0].keys()) |
|
t = Table(*table_headings, title="Drivers") |
|
|
|
for d in driver_dicts: |
|
# row_dict: dict[str, str] = {} |
|
row = [] |
|
|
|
for name, value in d.items(): |
|
if value is None: |
|
value = "" |
|
elif isinstance(value, MutableSequence): |
|
value = ",".join(value) |
|
elif isinstance(value, bool): |
|
value = "✅" if value else ":x:" |
|
elif isinstance(value, str) and name.endswith("url"): |
|
value = f"[link={value}]link[/link]" |
|
else: |
|
value = str(value) |
|
row.append(value) |
|
|
|
t.add_row(*row) |
|
c.print(t) |
|
|
|
|
|
if __name__ == "__main__": |
|
_APP() |