Merge pull request #3 from arangoml/jsonl_support

adds jsonl support

Author: cw00dw0rd

arango_datasets/datasets.py

@@ -1,5 +1,6 @@
+import json
 import sys
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List
 
 import requests
 from arango.collection import StandardCollection
@@ -25,13 +26,14 @@ class Datasets:
     def __init__(
         self,
         db: Database,
-        batch_size: Optional[int] = None,
+        batch_size: int = 50,
         metadata_file: str = "https://arangodb-dataset-library.s3.amazonaws.com/root_metadata.json",  # noqa: E501
     ):
         self.metadata_file: str = metadata_file
         self.metadata_contents: Dict[str, Any]
         self.batch_size = batch_size
         self.user_db = db
+        self.file_type: str
         if issubclass(type(db), Database) is False:
             msg = "**db** parameter must inherit from arango.database.Database"
             raise TypeError(msg)
@@ -74,36 +76,80 @@ def insert_docs(
         except DocumentInsertError as exec:
             print("Document insertion failed due to the following error:")
             print(exec.message)
+            sys.exit(1)
 
         print(f"Finished loading current file for collection: {collection_name}")
 
-    def load_file(self, collection_name: str, edge_type: bool, file_url: str) -> None:
-        collection: StandardCollection
-
+    def load_json(
+        self,
+        collection_name: str,
+        edge_type: bool,
+        file_url: str,
+        collection: StandardCollection,
+    ) -> None:
         try:
-            collection = self.user_db.create_collection(collection_name, edge=edge_type)
-        except CollectionCreateError as exec:
-            print(
-                f"""Failed to create {collection_name} collection due
-                 to the following error:"""
-            )
-            print(exec.error_message)
-            sys.exit(1)
+            with progress(f"Downloading file for: {collection_name}") as p:
+                p.add_task("load_file")
+                data = requests.get(file_url).json()
+        except (HTTPError, ConnectionError) as e:
+            print("Unable to download file.")
+            print(e)
+            raise e
+        print(f"Downloaded file for: {collection_name}, now importing... ")
+        self.insert_docs(collection, data, collection_name)
 
+    def load_jsonl(
+        self,
+        collection_name: str,
+        edge_type: bool,
+        file_url: str,
+        collection: StandardCollection,
+    ) -> None:
+        json_data = []
         try:
             with progress(f"Downloading file for: {collection_name}") as p:
                 p.add_task("load_file")
-                data = requests.get(file_url).json()
+                data = requests.get(file_url)
+
+            if data.encoding is None:
+                data.encoding = "utf-8"
+
+            for line in data.iter_lines(decode_unicode=True):
+                if line:
+                    json_data.append(json.loads(line))
+
         except (HTTPError, ConnectionError) as e:
             print("Unable to download file.")
             print(e)
             raise
         print(f"Downloaded file for: {collection_name}, now importing... ")
-        self.insert_docs(collection, data, collection_name)
+        self.insert_docs(collection, json_data, collection_name)
 
-    def load(self, dataset_name: str) -> None:
+    def load_file(self, collection_name: str, edge_type: bool, file_url: str) -> None:
+        collection: StandardCollection
+        try:
+            collection = self.user_db.create_collection(
+                collection_name, edge=edge_type
+            )  # type: ignore
+        except CollectionCreateError as exec:
+            print(
+                f"""Failed to create {collection_name} collection due
+                 to the following error:"""
+            )
+            print(exec.error_message)
+            sys.exit(1)
+        if self.file_type == "json":
+            self.load_json(collection_name, edge_type, file_url, collection)
+        elif self.file_type == "jsonl":
+            self.load_jsonl(collection_name, edge_type, file_url, collection)
+        else:
+            raise ValueError(f"Unsupported file type: {self.file_type}")
 
+    def load(self, dataset_name: str) -> None:
         if str(dataset_name).upper() in self.labels:
+            self.file_type = self.metadata_contents[str(dataset_name).upper()][
+                "file_type"
+            ]
 
             for edge in self.metadata_contents[str(dataset_name).upper()]["edges"]:
                 for e in edge["files"]:

tests/conftest.py

@@ -38,14 +38,14 @@ def pytest_configure(config: Any) -> None:
 
 def cleanup_collections() -> None:
     global db
-    if db.has_collection("flights"):
+    if db.has_collection("test_vertex"):
         try:
-            db.delete_collection("flights")
-        except (CollectionDeleteError):
-            print("unable to delete flights")
+            db.delete_collection("test_vertex")
+        except CollectionDeleteError:
+            print("unable to delete test_vertex")
 
-    if db.has_collection("airports"):
+    if db.has_collection("test_edge"):
         try:
-            db.delete_collection("airports")
-        except (CollectionDeleteError):
-            print("unable to delete airports")
+            db.delete_collection("test_edge")
+        except CollectionDeleteError:
+            print("unable to delete test_edge")

tests/test_main.py

@@ -1,12 +1,25 @@
 from typing import Any, no_type_check
 
 import pytest
+from requests import ConnectionError
 
 from arango_datasets.datasets import Datasets
 
-from .conftest import db
+from .conftest import cleanup_collections, db
 
+global test_metadata_url
+global root_metadata_url
+global bad_metadata_url
+test_metadata_url = (
+    "https://arangodb-dataset-library.s3.amazonaws.com/test_metadata.json"  # noqa: E501
+)
+root_metadata_url = (
+    "https://arangodb-dataset-library.s3.amazonaws.com/root_metadata.json"  # noqa: E501
+)
+bad_metadata_url = "http://bad_url.arangodb.com/"
 
+
+@no_type_check
 def test_dataset_constructor() -> None:
     assert Datasets(db)
     assert Datasets(db, batch_size=1000)
@@ -17,21 +30,31 @@ def test_dataset_constructor() -> None:
     assert Datasets(
         db,
         batch_size=1000,
-        metadata_file="https://arangodb-dataset-library.s3.amazonaws.com/root_metadata.json",  # noqa: E501
+        metadata_file=root_metadata_url,
     )
+    with pytest.raises(TypeError):
+        assert Datasets(
+            db="some none db object",
+            batch_size=1000,
+            metadata_file=root_metadata_url,
+        )
     with pytest.raises(Exception):
         assert Datasets({})
 
-    with pytest.raises(Exception):
-        assert Datasets(db, metadata_file="bad_url")
+    with pytest.raises(ConnectionError):
+        assert Datasets(db, metadata_file=bad_metadata_url)
 
 
+@no_type_check
 def test_list_datasets(capfd: Any) -> None:
-    datasets = Datasets(db).list_datasets()
+    datasets = Datasets(
+        db,
+        metadata_file=test_metadata_url,
+    ).list_datasets()
     out, err = capfd.readouterr()
-    assert "FLIGHTS" in out
+    assert "TEST" in out
     assert type(datasets) is list
-    assert "FLIGHTS" in datasets
+    assert "TEST" in datasets
 
 
 @no_type_check
@@ -42,17 +65,101 @@ def test_dataset_info(capfd: Any) -> None:
     with pytest.raises(Exception):
         Datasets(db).dataset_info(2)
 
-    dataset = Datasets(db).dataset_info("FLIGHTS")
+    dataset = Datasets(
+        db,
+        metadata_file=test_metadata_url,
+    ).dataset_info("TEST")
     assert type(dataset) is dict
 
+    assert dataset["TEST"]["file_type"] == "json"
+
     out, err = capfd.readouterr()
     assert len(out) > 0
 
 
+@no_type_check
+def test_load_file() -> None:
+    with pytest.raises(Exception):
+        Datasets.load_file(collection_name="test", edge_type=None, file_url="false")
+
+
+@no_type_check
+def test_load_json() -> None:
+    cleanup_collections()
+    collection_name = "test_vertex"
+    edge_type = False
+    file_url = "https://arangodb-dataset-library.s3.amazonaws.com/test_files/json/vertex_collection/test_vertex.json"  # noqa: E501
+    collection = db.create_collection("test_vertex")
+    assert None == (
+        Datasets.load_json(
+            Datasets(db),
+            collection_name=collection_name,
+            edge_type=edge_type,
+            file_url=file_url,
+            collection=collection,
+        )
+    )
+
+
+@no_type_check
+def json_bad_url() -> None:
+    cleanup_collections()
+    collection_name = "test_vertex"
+    edge_type = False
+    collection = db.create_collection("test_vertex")
+
+    with pytest.raises(ConnectionError):
+        Datasets.load_json(
+            Datasets(db),
+            collection_name=collection_name,
+            edge_type=edge_type,
+            file_url=bad_metadata_url,
+            collection=collection,
+        )
+
+
+@no_type_check
+def test_load_jsonl() -> None:
+    cleanup_collections()
+    collection_name = "test_vertex"
+    edge_type = False
+    file_url = "https://arangodb-dataset-library.s3.amazonaws.com/test_files/jsonl/vertex_collection/test_vertex.jsonl"  # noqa: E501
+    collection = db.create_collection("test_vertex")
+    assert None == (
+        Datasets.load_jsonl(
+            Datasets(db),
+            collection_name=collection_name,
+            edge_type=edge_type,
+            file_url=file_url,
+            collection=collection,
+        )
+    )
+
+
+@no_type_check
+def jsonl_bad_url() -> None:
+    cleanup_collections()
+    collection_name = "test_vertex"
+    edge_type = False
+    collection = db.create_collection("test_vertex")
+    with pytest.raises(ConnectionError):
+        Datasets.load_jsonl(
+            Datasets(db),
+            collection_name=collection_name,
+            edge_type=edge_type,
+            file_url=bad_metadata_url,
+            collection=collection,
+        )
+
+
 @no_type_check
 def test_load() -> None:
-    Datasets(db).load("FLIGHTS")
+    cleanup_collections()
+    Datasets(
+        db,
+        metadata_file=test_metadata_url,
+    ).load("TEST")
     with pytest.raises(Exception):
         Datasets(db).load(2)
-    assert db.collection("airports").count() == 3375
-    assert db.collection("flights").count() == 286463
+    assert db.collection("test_vertex").count() == 2
+    assert db.collection("test_edge").count() == 1