use prefix correctly (#7)

batchelor/reader.py

@@ -13,6 +13,14 @@ def parse_bucket(path: str) -> str:
     return path.split("/")[2]
 
 
+def parse_prefix(path: str) -> str:
+    """
+    Parse prefix from a GCS path. For example given the path
+    gs://bucket-name/path/to/file, return path/to/file
+    """
+    return path.split("/", 3)[3]
+
+
 def filter_json_files(paths: list[str]) -> list[str]:
     return [path for path in paths if path.endswith(".json") or path.endswith(".jsonl")]
 
@@ -22,7 +30,7 @@ def convert_path_to_list(path: str) -> list[str]:
         bucket_name = parse_bucket(path)
         paths = []
         client = storage.Client()
-        for blob in client.list_blobs(bucket_name, prefix=path):
+        for blob in client.list_blobs(bucket_name, prefix=parse_prefix(path)):
             paths.append(f"gs://{bucket_name}/{blob.name}")
         return filter_json_files(paths)
     return [path]

tests/test_reader.py

@@ -5,7 +5,7 @@
 from google.cloud import storage
 import pytest
 
-from batchelor.reader import parse_bucket, convert_path_to_list
+from batchelor.reader import parse_bucket, convert_path_to_list, parse_prefix
 
 
 @pytest.fixture
@@ -18,11 +18,21 @@ def test_parse_bucket(mock_client):
     expected = "bucket-name"
     assert parse_bucket(input) == expected
 
-    input = "gcss://bucket-name/path/to/file"
+    input = "gcs://bucket-name/path/to/file"
     expected = "bucket-name"
     assert parse_bucket(input) == expected
 
 
+def test_parse_bucket(mock_client):
+    input = "gs://bucket-name/path/to/file"
+    expected = "path/to/file"
+    assert parse_prefix(input) == expected
+
+    input = "gcs://bucket-name/path/to/file"
+    expected = "path/to/file"
+    assert parse_prefix(input) == expected
+
+
 @dataclass
 class Blob:
     name: str