Fix review comments

Author: deanp70

dagshub/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.6.10"
+__version__ = "0.6.9"
 from .logger import DAGsHubLogger, dagshub_logger
 from .common.init import init
 from .upload.wrapper import upload_files

dagshub/auth/token_auth.py

@@ -37,7 +37,7 @@ def auth_flow(self, request: Request) -> Generator[Request, Response, None]:
 
     def can_renegotiate(self):
         # Env var tokens cannot renegotiate, every other token type can
-        return type(self._token) is not EnvVarDagshubToken
+        return not type(self._token) is EnvVarDagshubToken
 
     def renegotiate_token(self):
         if not self._token_storage.is_valid_token(self._token, self._host):

dagshub/data_engine/annotation/importer.py

@@ -95,6 +95,7 @@ def download_annotations(self, dest_dir: Path):
         log_message("Downloading annotations from repository")
         repoApi = self.ds.source.repoApi
         if self.annotations_type == "cvat":
+            # Download just the annotation file
             repoApi.download(self.annotations_file.as_posix(), dest_dir, keep_source_prefix=True)
         elif self.annotations_type == "yolo":
             # Download the dataset .yaml file and the images + annotations
@@ -107,6 +108,7 @@ def download_annotations(self, dest_dir: Path):
             assert context.path is not None
             repoApi.download(self.annotations_file.parent / context.path, dest_dir, keep_source_prefix=True)
         elif self.annotations_type == "coco":
+            # Download just the annotation file
             repoApi.download(self.annotations_file.as_posix(), dest_dir, keep_source_prefix=True)
 
     @staticmethod
@@ -157,10 +159,8 @@ def remap_annotations(
                 )
                 continue
             for ann in anns:
-                if ann.filename is not None:
-                    ann.filename = remap_func(ann.filename)
-                else:
-                    ann.filename = new_filename
+                assert ann.filename is not None
+                ann.filename = remap_func(ann.filename)
             remapped[new_filename] = anns
 
         return remapped

dagshub/data_engine/annotation/metadata.py

@@ -271,28 +271,6 @@ def add_image_pose(
         self.annotations.append(ann)
         self._update_datapoint()
 
-    def add_coco_annotation(
-        self,
-        coco_json: str,
-    ):
-        """
-        Add annotations from a COCO-format JSON string.
-
-        Args:
-            coco_json: A COCO-format JSON string with ``categories``, ``images``, and ``annotations`` keys.
-        """
-        from dagshub_annotation_converter.converters.coco import load_coco_from_json_string
-
-        grouped, _ = load_coco_from_json_string(coco_json)
-        new_anns: list[IRAnnotationBase] = []
-        for anns in grouped.values():
-            for ann in anns:
-                ann.filename = self.datapoint.path
-                new_anns.append(ann)
-        self.annotations.extend(new_anns)
-        log_message(f"Added {len(new_anns)} COCO annotation(s) to datapoint {self.datapoint.path}")
-        self._update_datapoint()
-
     def add_yolo_annotation(
         self,
         annotation_type: Literal["bbox", "segmentation", "pose"],

dagshub/data_engine/model/query_result.py

@@ -900,7 +900,10 @@ def export_as_coco(
 
         context = CocoContext()
         if classes is not None:
-            context.categories = dict(classes)
+            categories = Categories()
+            for category_id, category_name in classes.items():
+                categories.add(category_name, category_id)
+            context.categories = categories
 
         # Add the source prefix to all annotations
         for ann in annotations:

tests/data_engine/annotation_import/test_coco.py

@@ -1,7 +1,6 @@
 import datetime
 import json
-from pathlib import PurePosixPath
-from unittest.mock import patch, PropertyMock
+from unittest.mock import patch
 
 import pytest
 from dagshub_annotation_converter.ir.image import (
@@ -17,12 +16,6 @@
 from dagshub.data_engine.model.query_result import QueryResult
 
 
-@pytest.fixture(autouse=True)
-def mock_source_prefix(ds):
-    with patch.object(type(ds.source), "source_prefix", new_callable=PropertyMock, return_value=PurePosixPath()):
-        yield
-
-
 # --- import ---
 
 
@@ -58,19 +51,6 @@ def test_coco_convert_to_ls_tasks(ds, tmp_path, mock_dagshub_auth):
     assert len(task_json["annotations"]) > 0
 
 
-# --- add_coco_annotation ---
-
-
-def test_add_coco_annotation_rewrites_filename(ds, mock_dagshub_auth):
-    dp = Datapoint(datasource=ds, path="my_images/photo.jpg", datapoint_id=0, metadata={})
-    meta_ann = MetadataAnnotations(datapoint=dp, field="ann")
-    meta_ann.add_coco_annotation(json.dumps(_make_coco_json()))
-
-    assert len(meta_ann.annotations) == 1
-    assert isinstance(meta_ann.annotations[0], IRBBoxImageAnnotation)
-    assert meta_ann.annotations[0].filename == "my_images/photo.jpg"
-
-
 # --- _resolve_annotation_field ---
 
 

tests/data_engine/conftest.py

@@ -1,4 +1,6 @@
 import datetime
+from pathlib import PurePosixPath
+from unittest.mock import PropertyMock
 
 import pytest
 
@@ -34,6 +36,7 @@ def _create_mock_datasource(mocker, id, name) -> Datasource:
     mocker.patch.object(ds_state, "get_from_dagshub")
     # Stub out root path so all the content_path/etc work without also mocking out RepoAPI
     mocker.patch.object(ds_state, "_root_path", return_value="http://example.com")
+    mocker.patch.object(type(ds_state), "source_prefix", new_callable=PropertyMock, return_value=PurePosixPath())
     ds_state.repoApi = MockRepoAPI("kirill/repo")
     return Datasource(ds_state)