diff options
Diffstat (limited to 'python/fatcat_tools/references.py')
-rw-r--r-- | python/fatcat_tools/references.py | 48 |
1 files changed, 37 insertions, 11 deletions
diff --git a/python/fatcat_tools/references.py b/python/fatcat_tools/references.py index 6fd9ca49..624020b5 100644 --- a/python/fatcat_tools/references.py +++ b/python/fatcat_tools/references.py @@ -124,7 +124,33 @@ class RefHits(BaseModel): limit: int query_time_ms: int query_wall_time_ms: int - result_refs: List[Union[BiblioRef, EnrichedBiblioRef]] + result_refs: List[BiblioRef] + + class Config: + json_encoders = { + ReleaseEntity: entity_to_dict, + } + + def as_enriched(self, enriched_refs: List[EnrichedBiblioRef]) -> "RefHitsEnriched": + return RefHitsEnriched( + count_returned=self.count_returned, + count_total=self.count_total, + offset=self.offset, + limit=self.limit, + query_time_ms=self.query_time_ms, + query_wall_time_ms=self.query_wall_time_ms, + result_refs=enriched_refs, + ) + + +class RefHitsEnriched(BaseModel): + count_returned: int + count_total: int + offset: int + limit: int + query_time_ms: int + query_wall_time_ms: int + result_refs: List[EnrichedBiblioRef] class Config: json_encoders = { @@ -221,7 +247,7 @@ def get_inbound_refs( limit: int = 25, offset: Optional[int] = None, es_index: str = "fatcat_ref", -) -> List[BiblioRef]: +) -> RefHits: search = Search(using=es_client, index=es_index) @@ -398,16 +424,16 @@ def run_ref_query(args) -> None: enriched = enrich_outbound_refs( hits.result_refs, hide="refs,abstracts", fatcat_api_client=args.fatcat_api_client ) - for ref in enriched: - if ref.release: + for eref in enriched: + if eref.release: print( - f"{ref.ref.ref_index or '-'}\trelease_{ref.release.ident}\t{ref.ref.match_provenance}/{ref.ref.match_status}\t{ref.release.release_year or '-'}\t{ref.release.title}\t{ref.release.ext_ids.pmid or ref.release.ext_ids.doi or '-'}" + f"{eref.ref.ref_index or '-'}\trelease_{eref.release.ident}\t{eref.ref.match_provenance}/{eref.ref.match_status}\t{eref.release.release_year or '-'}\t{eref.release.title}\t{eref.release.ext_ids.pmid or eref.release.ext_ids.doi or '-'}" ) else: - print(f"{ref.ref.ref_index or '-'}\trelease_{ref.target_release_ident}") + print(f"{eref.ref.ref_index or '-'}\trelease_{eref.ref.target_release_ident}") else: for ref in hits.result_refs: - print(f"{ref.ref.ref_index or '-'}\trelease_{ref.target_release_ident}") + print(f"{ref.ref_index or '-'}\trelease_{ref.target_release_ident}") print() print("## Inbound References") @@ -423,13 +449,13 @@ def run_ref_query(args) -> None: enriched = enrich_inbound_refs( hits.result_refs, hide="refs,abstracts", fatcat_api_client=args.fatcat_api_client ) - for ref in enriched: - if ref.release: + for eref in enriched: + if eref.release: print( - f"release_{ref.release.ident}\t{ref.ref.match_provenance}/{ref.ref.match_status}\t{ref.release.release_year or '-'}\t{ref.release.title}\t{ref.release.ext_ids.pmid or ref.release.ext_ids.doi or '-'}" + f"release_{eref.release.ident}\t{eref.ref.match_provenance}/{eref.ref.match_status}\t{eref.release.release_year or '-'}\t{eref.release.title}\t{eref.release.ext_ids.pmid or eref.release.ext_ids.doi or '-'}" ) else: - print(f"release_{ref.target_release_ident}") + print(f"release_{eref.ref.target_release_ident}") else: for ref in hits.result_refs: print(f"work_{ref.source_work_ident}\trelease_{ref.source_release_ident}") |