datatorch 0.6.2__tar.gz → 0.6.4__tar.gz

{datatorch-0.6.2 → datatorch-0.6.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datatorch
-Version: 0.6.2
+Version: 0.6.4
 Summary: A CLI and library for interacting with DataTorch.
 Home-page: https://github.com/datatorch/python
 Author: DataTorch

{datatorch-0.6.2 → datatorch-0.6.4}/datatorch/agent/pipelines/runner/runner.py

@@ -14,6 +14,12 @@ class ProcessCodeError(Exception):
     pass
 
 
+# Max size of one stdout line from an action subprocess. Output values
+# arrive as one `::var::<json>` line, so this bounds the largest single
+# output value (e.g. an http-request body), not ordinary log lines.
+STDOUT_LINE_LIMIT = 16 * 1024 * 1024
+
+
 class Runner(object):
     def __init__(self, config: dict, action: "Action"):
         self.config = config
@@ -110,6 +116,11 @@ class Runner(object):
             stdout=asyncio.subprocess.PIPE,
             stderr=asyncio.subprocess.STDOUT,
             env={**os.environ, **(env or {})},
+            # Outputs are reported as single `::var::<json>` stdout lines
+            # (see check_for_output), so a line can carry an entire fetched
+            # page body — asyncio's 64KB default readline limit kills the
+            # stream with LimitOverrunError on anything larger.
+            limit=STDOUT_LINE_LIMIT,
         )
         if wait:
             await process.wait()

{datatorch-0.6.2 → datatorch-0.6.4}/datatorch/api/api.py

@@ -56,6 +56,14 @@ _FILE = File.add_fragment("""
     }
     """)
 
+_DATASET = Dataset.add_fragment("""
+    query GetDataset($id: ID!) {
+      dataset: datasetById(id: $id) {
+        ...DatasetFields
+      }
+    }
+    """)
+
 _VIEWER = User.add_fragment("""
     query GetViewer {
       viewer {
@@ -119,6 +127,13 @@ class ApiClient(Client):
             File, self.query_to_class(File, _FILE, path="file", params={"fileId": id})
         )
 
+    def dataset(self, id: str) -> Dataset:
+        """Retrieve a dataset (or branch) by ID"""
+        return cast(
+            Dataset,
+            self.query_to_class(Dataset, _DATASET, path="dataset", params={"id": id}),
+        )
+
     def upload_to_default_filesource(
         self,
         project: Project,

{datatorch-0.6.2 → datatorch-0.6.4}/datatorch/api/client.py

@@ -162,7 +162,8 @@ class Client(object):
         dl_file = os.path.join(directory, name)
         dl_file = os.path.abspath(dl_file)
 
-        if os.path.isfile(dl_file) & skip == True:
+        # `isfile(x) & skip == True` only worked by bitwise-bool accident
+        if skip and os.path.isfile(dl_file):
             return dl_file, None
 
         query_string = urlencode({"download": "true", "stream": "true"})
@@ -172,8 +173,10 @@ class Client(object):
         result = requests.get(
             download_url, headers={self.token_header: self._api_token}, stream=True
         )
+        # Never write an error body to disk as a "downloaded file"
+        result.raise_for_status()
 
-        content = result.headers["content-disposition"]
+        content = result.headers.get("content-disposition", "")
         # _, value = cgi.parse_header(content)
         parser = HeaderParser()
         headers = parser.parsestr(f"Content-Disposition: {content}")

datatorch-0.6.4/datatorch/api/entity/dataset.py

@@ -0,0 +1,103 @@
+from .base import BaseEntity
+
+__all__ = "Dataset"
+
+
+_FORK_DATASET = """
+    mutation ForkDataset(
+      $datasetId: ID!
+      $branchName: String!
+      $versionId: ID
+    ) {
+      fork: forkDataset(
+        datasetId: $datasetId
+        branchName: $branchName
+        versionId: $versionId
+      ) {
+        dataset {
+          id
+          name
+          description
+          projectId
+          branchName
+          parentDatasetId
+          forkedAt
+        }
+      }
+    }
+"""
+
+
+_CREATE_DATASET = """
+    mutation CreateDataset(
+      $projectId: ID!
+      $name: String!
+      $description: String
+    ) {
+      dataset: createDataset(
+        input: {
+          projectId: $projectId
+          name: $name
+          description: $description
+        }
+      ) {
+        id
+      }
+    }
+"""
+
+
+class Dataset(BaseEntity):
+    id: str
+    name: str
+    description: str
+    project_id: str
+    kilobytes: int
+    formatted_bytes: int
+    created_at: str
+    updated_at: str
+
+    def create(self, client=None):
+        super().create(client=client)
+
+        assert self.project_id is not None
+        results = self.client.execute(
+            _CREATE_DATASET,
+            params={
+                "projectId": self.project_id,
+                "name": self.name,
+                "description": self.description,
+            },
+        )
+
+        self.id = results.get("dataset").get("id")
+
+    def fork(self, branch_name: str, version_id: str = None, client=None):
+        """
+        Forks this dataset into a branch and returns it as a new Dataset.
+
+        The branch shares the parent's storage objects and copies its
+        annotations. Pass `version_id` to start the branch from a version's
+        reconstructed state instead of the live state. The returned dataset
+        can be used anywhere a dataset id is accepted (e.g. uploads).
+
+        Note: branch attributes (`branch_name`, `parent_dataset_id`,
+        `forked_at`) are set dynamically on the returned instance; they are
+        deliberately kept out of the class annotations so the generated
+        `DatasetFields` fragment stays compatible with servers that predate
+        branching.
+        """
+        assert self.id is not None
+        if client:
+            self.client = client
+        assert self.client is not None
+
+        results = self.client.execute(
+            _FORK_DATASET,
+            params={
+                "datasetId": self.id,
+                "branchName": branch_name,
+                "versionId": version_id,
+            },
+        )
+        return Dataset(results.get("fork").get("dataset"), client=self.client)