Include law data in article search results

nfelger · Oct 12, 2020 · 3f20369 · 3f20369
1 parent 7d585d5
commit 3f20369
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 29 deletions.
diff --git a/rip_api/api.py b/rip_api/api.py
@@ -155,7 +155,7 @@ def get_search_results(
 ):
     orm_type_to_schema = {
         models.Law: api_schemas.LawBasicFields,
-        models.ContentItem: api_schemas.ContentItemBasicFields
+        models.ContentItem: api_schemas.ContentItemBasicFieldsWithLaw
     }
     with db.session_scope() as session:
         type_filter_value = type_filter and type_filter.value

diff --git a/rip_api/api_schemas.py b/rip_api/api_schemas.py
@@ -80,30 +80,6 @@ class ContentItemWithTextContent(ContentItemAllFields):
     documentaryFootnotes: Optional[str] = ...
 
 
-class HeadingBasicFields(ContentItemBasicFields):
-    type: str = "heading"
-
-
-class HeadingAllFields(ContentItemAllFields):
-    type: str = "heading"
-
-
-class ArticleBasicFields(ContentItemBasicFields):
-    type: str = "article"
-
-
-class ArticleAllFields(ContentItemWithTextContent):
-    type: str = "article"
-
-
-class HeadingArticleBasicFields(ContentItemBasicFields):
-    type: str = "headingArticle"
-
-
-class HeadingArticleAllFields(ContentItemWithTextContent):
-    type: str = "headingArticle"
-
-
 class LawBasicFields(pydantic.BaseModel):
     type: str = "law"
     id: str
@@ -175,6 +151,61 @@ def _attrs_dict_from_law(cls, law):
         return attrs
 
 
+class ContentItemBasicFieldsWithLaw(ContentItemBasicFields):
+    law: LawBasicFields
+
+    @staticmethod
+    def model_class_from_item_type(item_type):
+        return {
+            "article": ArticleBasicFieldsWithLaw,
+            "heading": HeadingBasicFieldsWithLaw,
+            "heading_article": HeadingArticleBasicFieldsWithLaw,
+        }[item_type]
+
+    @classmethod
+    def _attrs_dict_from_item(cls, item):
+        return {
+            **super()._attrs_dict_from_item(item),
+            "law": LawBasicFields.from_orm_model(item.law)
+        }
+
+
+class HeadingBasicFields(ContentItemBasicFields):
+    type: str = "heading"
+
+
+class HeadingBasicFieldsWithLaw(ContentItemBasicFieldsWithLaw):
+    type: str = "heading"
+
+
+class HeadingAllFields(ContentItemAllFields):
+    type: str = "heading"
+
+
+class ArticleBasicFields(ContentItemBasicFields):
+    type: str = "article"
+
+
+class ArticleBasicFieldsWithLaw(ContentItemBasicFieldsWithLaw):
+    type: str = "article"
+
+
+class ArticleAllFields(ContentItemWithTextContent):
+    type: str = "article"
+
+
+class HeadingArticleBasicFields(ContentItemBasicFields):
+    type: str = "headingArticle"
+
+
+class HeadingArticleBasicFieldsWithLaw(ContentItemBasicFieldsWithLaw):
+    type: str = "headingArticle"
+
+
+class HeadingArticleAllFields(ContentItemWithTextContent):
+    type: str = "headingArticle"
+
+
 class LawAllFieldsWithContents(LawAllFields):
     # Ordering in the Union matters, cf. LawResponse.
     contents: List[Union[ArticleAllFields, HeadingAllFields, HeadingArticleAllFields]]
@@ -218,6 +249,6 @@ class ContentItemResponse(pydantic.BaseModel):
 
 
 class SearchResultsResponse(pydantic.BaseModel):
-    data: List[Union[LawBasicFields, ArticleBasicFields, HeadingArticleBasicFields]]
+    data: List[Union[LawBasicFields, ArticleBasicFieldsWithLaw, ArticleBasicFields, HeadingArticleBasicFieldsWithLaw, HeadingArticleBasicFields]]
     links: PaginationLinks
     pagination: Pagination
diff --git a/tests/test_api.py b/tests/test_api.py
@@ -262,7 +262,7 @@ def test_full_text_search(self, client, law):
         results = response.json()["data"]
         assert len(results) == 3
 
-        assert results[0] == {
+        law_response_dict = {
             "type": "law",
             "id": "BJNR055429995",
             "url": "https://api.rechtsinformationsportal.de/v1/laws/skaufg",
@@ -276,21 +276,24 @@ def test_full_text_search(self, client, law):
             "abbreviation": "SkAufG",
             "slug": "skaufg"
         }
+        assert results[0] == law_response_dict
 
         assert results[1] == {
             "type": "headingArticle",
             "id": "BJNR055429995BJNG000200305",
             "url": "https://api.rechtsinformationsportal.de/v1/laws/skaufg/articles/BJNR055429995BJNG000200305",
             "name": "Art 2",
-            "title": None
+            "title": None,
+            "law": law_response_dict
         }
 
         assert results[2] == {
             "type": "article",
             "id": "BJNR055429995BJNE000801310",
             "url": "https://api.rechtsinformationsportal.de/v1/laws/skaufg/articles/BJNR055429995BJNE000801310",
             "name": "§ 2",
-            "title": "Grenzübertritt, Einreise"
+            "title": "Grenzübertritt, Einreise",
+            "law": law_response_dict
         }
 
     def test_includes_pagination_envelope(self, client, law):