Add JSON file output option (#5)

Author: ka-sarthak

src/nomad_actions/actions/entries/activities.py

@@ -37,7 +37,7 @@ async def create_artifact_subdirectory(data: CreateArtifactSubdirectoryInput) ->
 async def search(data: SearchInput) -> list[str]:
     """
     Activity to perform NOMAD search based on the provided input data. The search
-    results are written to a file in the specified format (Parquet or CSV) in the
+    results are written to a file in the specified format (Parquet, CSV, or JSON) in the
     artifacts directory.
 
     Args:
@@ -46,24 +46,30 @@ async def search(data: SearchInput) -> list[str]:
     Returns:
         list[str]: List of the generated output file paths.
     """
-    from nomad.search import search
+    from nomad.search import search as nomad_search
 
-    from nomad_actions.actions.entries.utils import write_csv_file, write_parquet_file
+    from nomad_actions.actions.entries.utils import (
+        write_csv_file,
+        write_json_file,
+        write_parquet_file,
+    )
 
     logger = activity.logger
 
     if data.output_file_type == 'parquet':
         write_dataset_file = write_parquet_file
     elif data.output_file_type == 'csv':
         write_dataset_file = write_csv_file
+    elif data.output_file_type == 'json':
+        write_dataset_file = write_json_file
     else:
-        raise ValueError('Unsupported file format. Please use parquet or csv.')
+        raise ValueError('Unsupported file format. Please use parquet, csv, or json.')
 
     generated_file_paths = []
 
     # first query
     search_counter = 1
-    response = search(
+    response = nomad_search(
         user_id=data.user_id,
         owner=data.owner,
         query=data.query,
@@ -86,7 +92,7 @@ async def search(data: SearchInput) -> list[str]:
         # create a copy to preserve the original pagination settings
         pagination = data.pagination.model_copy()
         pagination.page_after_value = response.pagination.next_page_after_value
-        response = search(
+        response = nomad_search(
             user_id=data.user_id,
             owner=data.owner,
             query=data.query,
@@ -111,7 +117,7 @@ async def search(data: SearchInput) -> list[str]:
 @activity.defn
 async def consolidate_output_files(data: ConsolidateOutputFilesInput) -> str:
     """
-    Activity to consolidate multiple Parquet or CSV files into a single file.
+    Activity to consolidate multiple Parquet, CSV, or JSON files into a single file.
 
     Args:
         data (ConsolidateOutputFilesInput): Input data for consolidating files.

src/nomad_actions/actions/entries/models.py

@@ -5,7 +5,7 @@
 from pydantic import BaseModel, Field
 
 OwnerLiteral = Literal['public', 'visible', 'shared', 'user', 'staging']
-OutputFileTypeLiteral = Literal['parquet', 'csv']
+OutputFileTypeLiteral = Literal['parquet', 'csv', 'json']
 IndexLiteral = Literal['entries', 'datasets', 'models', 'spaces']
 
 

src/nomad_actions/actions/entries/utils.py

@@ -44,8 +44,24 @@ def write_csv_file(path: str, data: list[dict]):
     df.to_csv(path, index=False, mode='w', header=True)
 
 
+def write_json_file(path: str, data: list[dict]):
+    """Writes a list of NOMAD entry dicts to a JSON file.
+
+    Args:
+        path (str): The path where the file will be saved.
+        data (list[dict]): The list of NOMAD entry dicts to be written to the file.
+    """
+    if not path.endswith('json'):
+        raise ValueError('Unsupported file type. Please use JSON.')
+
+    import json
+
+    with open(path, 'w') as f:
+        json.dump(data, f, indent=4)
+
+
 def consolidate_files(input_file_paths: list[str], output_file_path: str):
-    """Consolidates multiple Parquet or CSV files into a single file.
+    """Consolidates multiple Parquet, CSV, or JSON files into a single file.
 
     Args:
         input_file_paths (list[str]): List of file paths to be consolidated.
@@ -78,5 +94,16 @@ def consolidate_files(input_file_paths: list[str], output_file_path: str):
             dataframes.append(df)
         combined_df = pd.concat(dataframes, ignore_index=True)
         combined_df.to_csv(output_file_path, index=False)
+    elif output_file_path.endswith('json'):
+        import json
+
+        combined_data = []
+        for file_path in input_file_paths:
+            with open(file_path, encoding='utf-8') as f:
+                data = json.load(f)
+                # extend the combined entry list with entry list from each file
+                combined_data.extend(data)
+        with open(output_file_path, 'w', encoding='utf-8') as f:
+            json.dump(combined_data, f, indent=4)
     else:
-        raise ValueError('Unsupported file type. Please use parquet or CSV.')
+        raise ValueError('Unsupported file type. Please use parquet, CSV, or JSON.')