diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py index 3babc3872..f1a652924 100644 --- a/python/python/lancedb/util.py +++ b/python/python/lancedb/util.py @@ -34,20 +34,20 @@ def safe_import_adlfs(): adlfs = safe_import_adlfs() -def get_uri_scheme(uri: str) -> str: +def get_uri_scheme(uri: Union[str, pathlib.Path]) -> str: """ Get the scheme of a URI. If the URI does not have a scheme, assume it is a file URI. Parameters ---------- - uri : str + uri : str or pathlib.Path The URI to parse. Returns ------- str: The scheme of the URI. """ - parsed = urlparse(uri) + parsed = urlparse(str(uri)) scheme = parsed.scheme if not scheme: scheme = "file" @@ -61,25 +61,26 @@ def get_uri_scheme(uri: str) -> str: return scheme -def get_uri_location(uri: str) -> str: +def get_uri_location(uri: Union[str, pathlib.Path]) -> str: """ Get the location of a URI. If the parameter is not a url, assumes it is just a path Parameters ---------- - uri : str + uri : str or pathlib.Path The URI to parse. Returns ------- str: Location part of the URL, without scheme """ - parsed = urlparse(uri) + uri_str = str(uri) + parsed = urlparse(uri_str) if len(parsed.scheme) == 1: # Windows drive names are parsed as the scheme # e.g. "c:\path" -> ParseResult(scheme="c", netloc="", path="/path", ...) # So we add special handling here for schemes that are a single character - return uri + return uri_str if not parsed.netloc: return parsed.path diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index 741214794..4671109b8 100644 --- a/python/python/tests/test_util.py +++ b/python/python/tests/test_util.py @@ -30,6 +30,7 @@ import pytest import lancedb from lancedb.util import ( flatten_columns, + get_uri_location, get_uri_scheme, inf_vector_column_query, join_uri, @@ -86,6 +87,22 @@ def test_normalize_uri(): assert parsed_scheme == expected_scheme +def test_get_uri_scheme_and_location_pathlib(): + # Test pathlib.Path instances + abs_path = pathlib.Path("/tmp/data") + rel_path = pathlib.Path("relative/data") + + assert get_uri_scheme(abs_path) == "file" + assert get_uri_scheme(rel_path) == "file" + + assert get_uri_location(abs_path) == str(abs_path) + assert get_uri_location(rel_path) == str(rel_path) + + # String URIs still behave identically + assert get_uri_scheme("s3://bucket/data") == "s3" + assert get_uri_location("s3://bucket/data") == "bucket/data" + + def test_join_uri_remote(): schemes = ["s3", "az", "gs"] for scheme in schemes: