Merge pull request #128 from kurusugawa-computer/add-simple-annotation-parser

アノテーションzipをタスクごとにパースするメソッドを追加

Author: yuji38kwmt

annofabapi/__version__.py

@@ -1 +1 @@
-__version__ = '0.27.4'
+__version__ = '0.27.5'

annofabapi/parser.py

@@ -1,10 +1,11 @@
 import abc
 import json
 import os
+import re
 import warnings
 import zipfile
 from pathlib import Path
-from typing import Any, Iterator, List, Optional
+from typing import Any, Generic, Iterator, List, Optional, TypeVar
 
 from annofabapi.dataclass.annotation import FullAnnotation, SimpleAnnotation
 from annofabapi.exceptions import AnnotationOuterFileNotFoundError
@@ -292,6 +293,32 @@ def open_outer_file(self, data_uri: str):
             raise AnnotationOuterFileNotFoundError(str(outer_file_path))
 
 
+S = TypeVar("S", bound=SimpleAnnotationParser)
+
+
+class SimpleAnnotationParserGroupByTask(Generic[S]):
+    """
+    Simple Annotationのparserをタスクごとにまとめたもの。
+
+
+    Args:
+        task_id: タスクID
+        parser_list: タスク配下のJSONに関するパーサのList
+
+    """
+    def __init__(self, task_id: str, parser_list: List[S]):
+        self.__task_id = task_id
+        self.__parser_list = parser_list
+
+    @property
+    def task_id(self) -> str:
+        return self.__task_id
+
+    @property
+    def parser_list(self) -> List[S]:
+        return self.__parser_list
+
+
 def __parse_annotation_dir(annotaion_dir_path: Path, clazz) -> Iterator[Any]:
     for task_dir in annotaion_dir_path.iterdir():
         if not task_dir.is_dir():
@@ -333,6 +360,68 @@ def lazy_parse_full_annotation_dir(annotaion_dir_path: Path) -> Iterator[SimpleA
     return __parse_annotation_dir(annotaion_dir_path, FullAnnotationDirParser)
 
 
+def lazy_parse_simple_annotation_zip_by_task(
+        zip_file_path: Path) -> Iterator[SimpleAnnotationParserGroupByTask[SimpleAnnotationZipParser]]:
+    """
+    Simpleアノテーションzipファイル内を探索し、タスクごとに各annotationをparse可能なオブジェクトの列を返します。
+
+    Args:
+        zip_file_path: annofabからダウンロードしたsimple annotationのzipファイルへのパス
+
+    Yields:
+        対象タスク内の、annotationの遅延Parseが可能なインスタンス列
+    """
+    def is_input_data_info_in_task(zip_info: zipfile.ZipInfo, task_id: str) -> bool:
+        """
+        指定されたtask_id配下の入力データJSONかどうか
+        """
+        paths = [p for p in zip_info.filename.split("/") if len(p) != 0]
+        if len(paths) != 2:
+            return False
+        if paths[0] != task_id:
+            return False
+        if not paths[1].endswith(".json"):
+            return False
+
+        return True
+
+    with zipfile.ZipFile(zip_file_path, mode="r") as file:
+        info_list: List[zipfile.ZipInfo] = file.infolist()
+        # 1階層目のディレクトリをtask_idとみなす
+        task_info_list = [e for e in info_list if e.is_dir() and len(re.findall("/", e.filename)) == 1]
+
+        for task_info in task_info_list:
+            task_id = task_info.filename.split("/")[0]
+            parser_list = [
+                SimpleAnnotationZipParser(file, e.filename)
+                for e in info_list
+                if is_input_data_info_in_task(e, task_id)
+            ]
+
+            yield SimpleAnnotationParserGroupByTask(task_id, parser_list)
+
+
+def lazy_parse_simple_annotation_dir_by_task(
+        annotaion_dir_path: Path) -> Iterator[SimpleAnnotationParserGroupByTask[SimpleAnnotationDirParser]]:
+    """
+    Simpleアノテーションzipを展開したディレクトリ内を探索し、タスクごとに各annotationをparse可能なオブジェクトの列を返します。
+
+    Args:
+        annotaion_dir_path: annofabからダウンロードしたsimple annotationのzipファイルを展開したディレクトリ
+
+    Yields:
+        対象タスク内の、annotationの遅延Parseが可能なインスタンス列
+    """
+
+    for task_dir in annotaion_dir_path.iterdir():
+        if not task_dir.is_dir():
+            continue
+
+        task_id = task_dir.name
+        parser_list = [SimpleAnnotationDirParser(e) for e in task_dir.iterdir() if e.is_file() and e.suffix == ".json"]
+        yield SimpleAnnotationParserGroupByTask(task_id, parser_list)
+
+
 def __parse_annotation_zip(zip_file_path: Path, clazz) -> Iterator[Any]:
     def lazy_parser(zip_file: zipfile.ZipFile, info: zipfile.ZipInfo) -> Optional[Any]:
         paths = [p for p in info.filename.split("/") if len(p) != 0]

tests/test_local_parser.py

@@ -99,6 +99,34 @@ def test_simple_annotation_dir(self):
         with pytest.raises(AnnotationOuterFileNotFoundError):
             parser.open_outer_file("foo")
 
+    def test_lazy_parse_simple_annotation_zip_by_task(self):
+        zip_path = Path(test_dir / "simple-annotation-v2.zip")
+        task_parser_list = list(annofabapi.parser.lazy_parse_simple_annotation_zip_by_task(zip_path))
+
+        assert len(task_parser_list) == 2
+        assert len([e for e in task_parser_list if e.task_id == "sample_1"]) == 1
+        assert len([e for e in task_parser_list if e.task_id == "sample_0"]) == 1
+
+        task_parser = [e for e in task_parser_list if e.task_id == "sample_1"][0]
+        parser_list = task_parser.parser_list
+        assert len(parser_list) == 2
+        assert len([e for e in parser_list if e.input_data_id == "c6e1c2ec-6c7c-41c6-9639-4244c2ed2839"]) == 1
+        assert len([e for e in parser_list if e.input_data_id == "c86205d1-bdd4-4110-ae46-194e661d622b"]) == 1
+
+    def test_lazy_parse_simple_annotation_dir_by_task(self):
+        zip_path = Path(test_dir / "simple-annotation-v2")
+        task_parser_list = list(annofabapi.parser.lazy_parse_simple_annotation_dir_by_task(zip_path))
+
+        assert len(task_parser_list) == 2
+        assert len([e for e in task_parser_list if e.task_id == "sample_1"]) == 1
+        assert len([e for e in task_parser_list if e.task_id == "sample_0"]) == 1
+
+        task_parser = [e for e in task_parser_list if e.task_id == "sample_1"][0]
+        parser_list = task_parser.parser_list
+        assert len(parser_list) == 2
+        assert len([e for e in parser_list if e.input_data_id == "c6e1c2ec-6c7c-41c6-9639-4244c2ed2839"]) == 1
+        assert len([e for e in parser_list if e.input_data_id == "c86205d1-bdd4-4110-ae46-194e661d622b"]) == 1
+
 
 class TestFullAnnotation:
     def test_full_annotation_zip(self):