adding cli option, save_to_text and print flags

Author: ArturU043

pyproject.toml

@@ -36,5 +36,9 @@ test = [
     "pandas"
 ]
 
+[project.scripts]
+servicex-get-structure = "servicex_analysis_utils.cli:run_from_command"
+
+
 [tool.hatch.build.targets.wheel]
 packages = ["servicex_analysis_utils"]

servicex_analysis_utils/cli.py

@@ -0,0 +1,44 @@
+import argparse
+import sys
+import json
+import os
+from .file_peeking import get_structure 
+
+def run_from_command():
+    parser = argparse.ArgumentParser(description="CLI tool for retrieving ROOT file structures.")
+
+    parser.add_argument("dataset", nargs='+', help="Input datasets (Rucio DID) or a JSON file containing datasets in a dict.")  
+    parser.add_argument("--filter-branch", default="", help="Only display branches containing this string.")
+    parser.add_argument("--save-to-txt", action="store_true", help="Save output to a text file instead of printing.")
+
+    args = parser.parse_args()
+
+    if len(args.dataset) == 1 and args.dataset[0].endswith(".json"):
+        dataset_file = args.dataset[0]
+        
+        if not os.path.isfile(dataset_file):
+            print(f"\033[91mError: JSON file '{dataset_file}' not found.\033[0m", file=sys.stderr)
+            sys.exit(1)
+
+        try:
+            with open(dataset_file, "r") as f:
+                dataset = json.load(f) 
+                
+                if not isinstance(dataset, dict):
+                    print(f"\033[91mError: The JSON file must contain a dictionary.\033[0m", file=sys.stderr)
+                    sys.exit(1)
+            
+        except json.JSONDecodeError:
+            print(f"\033[91mError: '{dataset_file}' is not a valid JSON file.\033[0m", file=sys.stderr)
+            sys.exit(1)
+
+    else:
+        # If dataset is provided directly in CLI, use it as a list
+        dataset = args.dataset
+
+    result = get_structure(dataset, filter_branch=args.filter_branch, save_to_txt=args.save_to_txt, do_print=False)
+
+    if not args.save_to_txt:
+        print(result)
+    else:
+        print("Saved to samples_structure.txt")

servicex_analysis_utils/file_peeking.py

@@ -86,62 +86,61 @@ def is_tree(obj):
     # Return str in an array
     return ak.Array([final_str])
 
-
-def print_structure_from_str(deliver_dict, filter_branch="", save_to_txt=False ):
+def print_structure_from_str(deliver_dict, filter_branch="", save_to_txt=False, do_print=False):
     """
-    Helper. Takes the structure strings for all samples from servicex.deliver output 
-    and prints them in a friendly formatted view.
-    
-    The expected structure string format is:
-    
-      Tree: TreeName1; TBranch: Branchname1 ; dtype: BranchType1, TBranch: Branchname2 ; dtype: BranchType2, ...
-      Tree: TreeName2; TBranch: Branchname1 ; dtype: BranchType1, ...
-    
+    Converts dataset file structures to a formatted string.
+
     Parameters:
-      deliver_dict (dict): The return dictionary of servicex.deliver
-                          (keys are sample names, values are file paths or URLs)
-      filter_branch (str): Optional. Only Branch names containing it are printed.
-      save_to_txt (bool): Optional. Select if file structure is printed or dumped to .txt
+      deliver_dict (dict): ServiceX deliver output (keys: sample names, values: file paths or URLs).
+      filter_branch (str): If provided, only branches containing this string are included.
+      save_to_txt (bool): If True, saves output to a text file instead of returning it.
+
+    Returns:
+      str: The formatted file structure.
     """
-    print(f"File structure of all samples with branch filter {filter_branch}:")
+    output_lines = []
+    output_lines.append(f"\nFile structure of all samples with branch filter '{filter_branch}':")
 
     for sample_name, path in deliver_dict.items():
-        #Sample name with icon and bands
-        print(
+        output_lines.append(
             f"\n---------------------------\n"
             f"\U0001F4C1 Sample: {sample_name}\n"
             f"---------------------------"
         )
 
         with uproot.open(path[0]) as f:
-            #Expected position of structure_str from servicex.deliver
-            structure_str=f["servicex"]["branch"].array()[0]
-        
-        # Split at each \n each new line represents one tree.
+            structure_str = f["servicex"]["branch"].array()[0]
+
         tree_lines = structure_str.split("\n")
-        
         for line in tree_lines:
             if not line.strip():
                 continue  # Skip empty lines
 
-            # First part before ';' is the tree header.
             parts = line.split(";", 1)
-            tree_header = parts[0] 
-            print(f"\n\U0001F333 {tree_header}")  # Print tree header with icon
-            
-            # Check for branches
+            tree_header = parts[0]
+            output_lines.append(f"\n\U0001F333 {tree_header}")
+
             if len(parts) > 1:
-                # branch info separated by ','
-                branch_info_str = parts[1]
-                branch_infos = branch_info_str.split(",")
-                print("   ├── Branches:")
+                branch_infos = parts[1].split(",")
+                output_lines.append("   ├── Branches:")
                 for b in branch_infos:
                     branch_line = b.strip()
                     if filter_branch not in branch_line:
                         continue
-                    # Only print lines that start with "TBranch:"
                     if branch_line.startswith("TBranch:"):
-                        print(f"   │   ├── {branch_line[8:]}")
+                        output_lines.append(f"   │   ├── {branch_line[8:]}")
+    
+    result_str = "\n".join(output_lines)
+
+    if save_to_txt:
+        with open("samples_structure.txt", "w") as f:
+            f.write(result_str)
+        return "File structure saved to 'samples_structure.txt'."
+    if do_print:
+        print(result_str)
+        return 
+    else:
+        return result_str
 
 
 def get_structure(dataset, **kwargs):
@@ -171,24 +170,21 @@ def get_structure(dataset, **kwargs):
         dataset_dict=dataset
     else:
         raise ValueError(f"Unsupported dataset input type: {user_in}.\nInput must be dict ('sample_name':'dataset_id'), str or list of str")
+        return 0
 
-    sample_list=[]
-
-    for name, did in dataset_dict.items():
-        tmp_dict={
-            "NFiles":1,
+    sample_list = [
+        {
+            "NFiles": 1,
             "Name": name,
             "Dataset": servicex.dataset.Rucio(did),
             "Query": query_PythonFunction,
         }
-        sample_list.append(tmp_dict)
+        for name, did in dataset_dict.items()
+    ]
 
-    spec_python = {
-    "Sample": sample_list
-    }
+    spec_python = {"Sample": sample_list}
 
     output=servicex.deliver(spec_python)
 
-    print_structure_from_str(output,**kwargs)
-    
+    return print_structure_from_str(output, **kwargs)