refactor: better logic of required columns defining

Author: andreishket

server/workers/common/common/utils.py

@@ -139,23 +139,30 @@ def push_metadata_to_queue(
     :param source_list: define from which service additional metadata will be received (available values: "crossref", "altmetric").
     :return: request_id for the receiving of the request result.
     """
+    # Checks that valid values are specified in the source array
     check_metadata_enrichment_source(source_list)
 
+    # Creates a new unique request identifier that will then be used to retrieve the result
     request_id = str(uuid.uuid4())
+
+    # Specifies from which sources to obtain information
     params["metrics_sources"] = source_list
+
+    # Payload object creation
     task_data = json.dumps({
         "id": request_id,
         "params": params,
         "metadata": metadata.to_json(orient="records"),
     })
 
+    # Pushing request to Redis and returning request id
     redis_store.rpush("metrics", task_data)
     return request_id
 
 
 def check_metadata_enrichment_source(source_list: List[str]) -> None:
     """
-    Checks that source for metadata enrichment contains correct values.
+    Checks that valid values are specified in the source array.
 
     :param source_list: List of sources from where metadata will be enriched.
     :return: None.
@@ -173,23 +180,29 @@ def fetch_enriched_metadata(redis_store: redis.Redis, request_id: str, timeout:
     :param timeout: Results waiting time (default - 600 seconds).
     :return: Enriched DataFrame with metadata.
     """
+    # Getting result of metadata enrichment from Redis
     result = get_key(redis_store, request_id, timeout)
     return pd.DataFrame(result["input_data"])
 
 
-def get_metadata_columns_for_integration(integration: Literal["pubmed", "orcid"]):
+def get_metadata_columns_for_source(source_list: List[str]) -> List[str]:
     """
-    Returning required metadata columns for different integrations.
+    Returning required metadata columns for different sources.
 
-    :param integration: integration service.
+    :param source_list: List of sources from where metadata received.
     :return: array with required metadata columns.
     """
+    # Checks that valid values are specified in the source array
+    check_metadata_enrichment_source(source_list)
+
+    # Define required metadata columns for different sources and return them
+    result = []
 
-    if integration == 'pubmed':
-        return ["citation_count"]
-    elif integration == 'orcid':
-        return [
-            "citation_count",
+    if "crossref" in source_list:
+        result.extend(["citation_count"])
+
+    if "altmetric" in source_list:
+        result.extend([
             "cited_by_wikipedia_count",
             "cited_by_msm_count",
             "cited_by_policies_count",
@@ -202,21 +215,24 @@ def get_metadata_columns_for_integration(integration: Literal["pubmed", "orcid"]
             "cited_by_qna_count",
             "cited_by_tweeters_count",
             "cited_by_videos_count"
-        ]
+        ])
 
-    return []
+    return result
 
 
-def ensure_required_columns(metadata: pd.DataFrame, integration: Literal["pubmed", "orcid"]) -> pd.DataFrame:
+def ensure_required_columns(metadata: pd.DataFrame, source_list: List[str]) -> pd.DataFrame:
     """
     Checks that all necessary columns are available or adding them with NaN value.
 
     :param metadata: DataFrame with metadata.
-    :param integration: integration service.
+    :param source_list: List of sources from where metadata received.
     :return: Updated DataFrame.
     """
+    # Checks that valid values are specified in the source array
+    check_metadata_enrichment_source(source_list)
 
-    columns = get_metadata_columns_for_integration(integration)
+    # Gets metadata columns that must be received from source(-s)
+    columns = get_metadata_columns_for_source(source_list)
     for column in columns:
         if column not in metadata.columns:
             metadata[column] = np.NaN
@@ -229,7 +245,6 @@ def enrich_metadata(
     params: Dict[str, Union[str, List[str]]],
     metadata: pd.DataFrame,
     source_list: List[str],
-    integration: Literal["pubmed", "orcid"]
 ) -> pd.DataFrame:
     """
     Enriching metadata - adding information about citations from Redis.
@@ -240,7 +255,7 @@ def enrich_metadata(
     :param source: define from which service additional metadata will be received (available values: "crossref", "altmetric").
     :return: Enriched DataFrame with metadata.
     """
-    # Checks that source list contains valid values
+    # Checks that valid values are specified in the source array
     check_metadata_enrichment_source(source_list)
 
     # Creates a request to metrics for metadata enrichment
@@ -251,5 +266,5 @@ def enrich_metadata(
     enriched_metadata = fetch_enriched_metadata(redis, request_id)
 
     # Checks that all necessary columns are available or adding them with NaN value
-    enriched_metadata = ensure_required_columns(enriched_metadata, integration)
+    enriched_metadata = ensure_required_columns(enriched_metadata, source_list)
     return enriched_metadata

server/workers/orcid/src/orcid_service.py

@@ -389,9 +389,8 @@ def _retrieve_author_info_and_metadata(self, orcid: Orcid) -> Tuple[AuthorInfo,
     def _process_metadata(self, metadata: pd.DataFrame, author_info: AuthorInfo, params: Dict[str, str]) -> pd.DataFrame:
         metadata["authors"] = metadata["authors"].replace("", author_info.author_name)
 
-        integration = 'orcid'
         source_for_metadata_enrichment = ["crossref", "altmetric"]
-        metadata = enrich_metadata(self.redis_store, params, metadata, source_for_metadata_enrichment, integration)
+        metadata = enrich_metadata(self.redis_store, params, metadata, source_for_metadata_enrichment)
 
         self.logger.debug(f'metadata shape after base enrichment: {metadata.shape}')
         author_info = self.enrich_author_info(author_info, metadata, params)

server/workers/pubmed/src/pubmed.py

@@ -42,9 +42,8 @@ def execute_search(self, params):
 
             metadata = pd.DataFrame(raw_metadata)
 
-            integration = 'pubmed'
             source_for_metadata_enrichment = ["crossref"]
-            metadata = enrich_metadata(self.redis_store, params, metadata, source_for_metadata_enrichment, integration)
+            metadata = enrich_metadata(self.redis_store, params, metadata, source_for_metadata_enrichment)
 
             text = pd.DataFrame(raw_text)