Fixed #107 #108

pepdbagent/_version.py

@@ -1 +1 @@
-__version__ = "0.7.0a2"
+__version__ = "0.7.0a3"

pepdbagent/modules/annotation.py

@@ -50,6 +50,7 @@ def get(
         filter_by: Optional[Literal["submission_date", "last_update_date"]] = None,
         filter_start_date: Optional[str] = None,
         filter_end_date: Optional[str] = None,
+        pep_type: Optional[Literal["pep", "pop"]] = None,
     ) -> AnnotationList:
         """
         Get project annotations.
@@ -77,6 +78,7 @@ def get(
             [Default: filter won't be used]
         :param filter_start_date: Filter start date. Format: "YYYY/MM/DD"
         :param filter_end_date: Filter end date. Format: "YYYY/MM/DD". if None: present date will be used
+        :param pep_type: Get pep with specified type. Options: ["pep", "pop"]. Default: None, get all peps
         :return: pydantic model: AnnotationList
         """
         if all([namespace, name, tag]):
@@ -94,6 +96,10 @@ def get(
                 offset=0,
                 results=found_annotation,
             )
+
+        if pep_type not in [None, "pep", "pop"]:
+            raise ValueError(f"pep_type should be one of ['pep', 'pop'], got {pep_type}")
+
         return AnnotationList(
             limit=limit,
             offset=offset,
@@ -104,6 +110,7 @@ def get(
                 filter_by=filter_by,
                 filter_end_date=filter_end_date,
                 filter_start_date=filter_start_date,
+                pep_type=pep_type,
             ),
             results=self._get_projects(
                 namespace=namespace,
@@ -116,6 +123,7 @@ def get(
                 filter_by=filter_by,
                 filter_end_date=filter_end_date,
                 filter_start_date=filter_start_date,
+                pep_type=pep_type,
             ),
         )
 
@@ -222,9 +230,11 @@ def _count_projects(
         filter_by: Optional[Literal["submission_date", "last_update_date"]] = None,
         filter_start_date: Optional[str] = None,
         filter_end_date: Optional[str] = None,
+        pep_type: Optional[Literal["pep", "pop"]] = None,
     ) -> int:
         """
         Count projects. [This function is related to _find_projects]
+
         :param namespace: namespace where to search for a project
         :param search_str: search string. will be searched in name, tag and description information
         :param admin: string or list of admins [e.g. "Khoroshevskyi", or ["doc_adin","Khoroshevskyi"]]
@@ -233,6 +243,8 @@ def _count_projects(
             [Default: filter won't be used]
         :param filter_start_date: Filter start date. Format: "YYYY:MM:DD"
         :param filter_end_date: Filter end date. Format: "YYYY:MM:DD". if None: present date will be used
+        :param pep_type: Get pep with specified type. Options: ["pep", "pop"]. Default: None, get all peps
+
         :return: number of found project in specified namespace
         """
         if admin is None:
@@ -247,6 +259,8 @@ def _count_projects(
         statement = self._add_date_filter_if_provided(
             statement, filter_by, filter_start_date, filter_end_date
         )
+        if pep_type:
+            statement = statement.where(Projects.pop.is_(pep_type == "pop"))
         result = self._pep_db_engine.session_execute(statement).first()
 
         try:
@@ -266,6 +280,7 @@ def _get_projects(
         filter_by: Optional[Literal["submission_date", "last_update_date"]] = None,
         filter_start_date: Optional[str] = None,
         filter_end_date: Optional[str] = None,
+        pep_type: Optional[Literal["pep", "pop"]] = None,
     ) -> List[AnnotationModel]:
         """
         Get projects by providing search string.
@@ -284,6 +299,7 @@ def _get_projects(
             [Default: filter won't be used]
         :param filter_start_date: Filter start date. Format: "YYYY:MM:DD"
         :param filter_end_date: Filter end date. Format: "YYYY:MM:DD". if None: present date will be used
+        :param pep_type: Get pep with specified type. Options: ["pep", "pop"]. Default: None, get all peps
         :return: list of found projects with their annotations.
         """
         _LOGGER.info(f"Running annotation search: (namespace: {namespace}, query: {search_str}.")
@@ -303,6 +319,8 @@ def _get_projects(
         )
         statement = self._add_order_by_keyword(statement, by=order_by, desc=order_desc)
         statement = statement.limit(limit).offset(offset)
+        if pep_type:
+            statement = statement.where(Projects.pop.is_(pep_type == "pop"))
 
         id_results = self._pep_db_engine.session_execute(statement).all()
 
@@ -463,3 +481,82 @@ def get_project_number_in_namespace(
             return result[0]
         except IndexError:
             return 0
+
+    def get_by_rp_list(
+        self,
+        registry_paths: List[str],
+        admin: Union[str, List[str]] = None,
+    ) -> AnnotationList:
+        """
+        Get project annotations by providing list of registry paths.
+
+        :param registry_paths: registry path string or list of registry paths
+        :param admin: list of namespaces where user is admin
+        :return: pydantic model: AnnotationReturnModel(
+            limit:
+            offset:
+            count:
+            result: List [AnnotationModel])
+        """
+        admin_tuple = tuple_converter(admin)
+
+        if isinstance(registry_paths, list):
+            or_statement_list = []
+            for path in registry_paths:
+                try:
+                    namespace, name, tag = registry_path_converter(path)
+                    or_statement_list.append(
+                        and_(
+                            Projects.name == name,
+                            Projects.namespace == namespace,
+                            Projects.tag == tag,
+                            or_(
+                                Projects.namespace.in_(admin_tuple),
+                                Projects.private.is_(False),
+                            ),
+                        )
+                    )
+                except RegistryPathError as err:
+                    _LOGGER.error(str(err), registry_paths)
+                    continue
+            if not or_statement_list:
+                _LOGGER.error("No valid registry paths were provided!")
+                return AnnotationList(
+                    count=0,
+                    limit=len(registry_paths),
+                    offset=0,
+                    results=[],
+                )
+
+            statement = select(Projects).where(or_(*or_statement_list))
+            anno_results = []
+            with Session(self._sa_engine) as session:
+                query_result = session.execute(statement).all()
+                for result in query_result:
+                    project_obj = result[0]
+                    annot = AnnotationModel(
+                        namespace=project_obj.namespace,
+                        name=project_obj.name,
+                        tag=project_obj.tag,
+                        is_private=project_obj.private,
+                        description=project_obj.description,
+                        number_of_samples=project_obj.number_of_samples,
+                        submission_date=str(project_obj.submission_date),
+                        last_update_date=str(project_obj.last_update_date),
+                        digest=project_obj.digest,
+                        pep_schema=project_obj.pep_schema,
+                        pop=project_obj.pop,
+                        stars_number=len(project_obj.stars_mapping),
+                    )
+                    anno_results.append(annot)
+
+            return_len = len(anno_results)
+            return AnnotationList(
+                count=return_len,
+                limit=len(registry_paths),
+                offset=0,
+                results=anno_results,
+            )
+
+        else:
+            return self.get_by_rp(registry_paths, admin)

tests/test_pepagent.py

@@ -508,6 +508,51 @@ def test_search_incorrect_filter_by_string(
                 filter_end_date=date_now.strftime("%Y/%m/%d"),
             )
 
+    @pytest.mark.parametrize(
+        "rp_list, admin, found_number",
+        [
+            [
+                [
+                    "namespace1/amendments1:default",
+                    "namespace1/amendments2:default",
+                    "namespace2/derive:default",
+                    "private_test/amendments1:default",
+                ],
+                "namespace1",
+                3,
+            ],
+            [
+                [
+                    "namespace1/amendments1:default",
+                    "namespace1/amendments2:default",
+                    "namespace2/derive:default",
+                    "private_test/amendments1:default",
+                ],
+                "private_test",
+                4,
+            ],
+        ],
+    )
+    def test_get_annotation_by_rp_list(self, initiate_pepdb_con, rp_list, admin, found_number):
+        result = initiate_pepdb_con.annotation.get_by_rp_list(rp_list)
+        assert len(result.results) == 3
+
+    def test_get_annotation_by_rp_enpty_list(self, initiate_pepdb_con):
+        result = initiate_pepdb_con.annotation.get_by_rp_list([])
+        assert len(result.results) == 0
+
+    @pytest.mark.parametrize(
+        "namespace, query, found_number",
+        [
+            ["namespace1", "ame", 2],
+        ],
+    )
+    def test_search_incorrect_incorrect_pep_type(
+        self, initiate_pepdb_con, namespace, query, found_number
+    ):
+        with pytest.raises(ValueError):
+            initiate_pepdb_con.annotation.get(namespace=namespace, pep_type="incorrect")
+
 
 @pytest.mark.skipif(
     not db_setup(),