Get DS file structure with serviceX tool

.github/workflows/CI.yml

@@ -6,7 +6,26 @@ on:
   workflow_dispatch:
 
 jobs:
+
+  black-format:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.x"
+
+      - name: Run Black
+        run: |
+          pipx run black --check .
+
   test:
+    needs:
+      - black-format
+
     runs-on: ubuntu-latest
 
     steps:

.gitignore

@@ -7,3 +7,10 @@ servicex.yaml
 
 #Distribution 
 dist/
+
+#ServiceX
+servicex.yaml
+
+#Testing 
+samples_structure.txt
+

pyproject.toml

@@ -11,6 +11,7 @@ readme = "README.md"
 license = { text = "BSD-3-Clause" }
 requires-python = ">=3.9"
 dependencies = [ 
+    "servicex",
     "uproot>=5.0",
     "awkward>=2.6",
     "dask-awkward>=2024.12.2",
@@ -32,8 +33,13 @@ test = [
     "pytest>=7.2.0",
     "numpy>=1.21", 
     "pyarrow>=8.0.0",
-    "pandas"
+    "pandas",
+    "miniopy-async==1.21.1"
 ]
 
+[project.scripts]
+servicex-get-structure = "servicex_analysis_utils.cli:app"
+
+
 [tool.hatch.build.targets.wheel]
 packages = ["servicex_analysis_utils"]

servicex_analysis_utils/__init__.py

@@ -25,7 +25,8 @@
 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-from .materialization import to_awk 
+from .materialization import to_awk
+from .file_peeking import get_structure
 
 __version__ = "1.0.0"
-__all__ = ['to_awk']
+__all__ = ["to_awk", "get_structure"]

servicex_analysis_utils/cli.py

@@ -0,0 +1,73 @@
+import sys
+import json
+import os
+import logging
+from .file_peeking import get_structure
+import typer
+from typing import List
+
+app = typer.Typer()
+
+
+def make_dataset_list(dataset_arg):
+    """
+    Helper to handle the user input daset argument.
+    Loads to dict if input is .json else returns default input
+    Output is given to get_structure()
+
+    Parameters:
+    dataset_arg (str, [str]): Single DS identifier, list of multiple identifiers or path/to/.json containig identifiers and sample names.
+
+    Returns:
+    dataset (str, [str], dict): dictionary loaded from the json
+    """
+    if len(dataset_arg) == 1 and dataset_arg[0].endswith(".json"):
+        dataset_file = dataset_arg[0]
+
+        if not os.path.isfile(dataset_file):
+            logging.error(f"Error: JSON file '{dataset_file}' not found.")
+            sys.exit(1)
+
+        try:
+            with open(dataset_file, "r") as f:
+                dataset = json.load(f)
+
+                if not isinstance(dataset, dict):
+                    logging.error(f"Error: The JSON file must contain a dictionary.")
+                    sys.exit(1)
+
+        except json.JSONDecodeError:
+            logging.error(
+                f"Error: '{dataset_file}' is not a valid JSON file.", exc_info=True
+            )
+            sys.exit(1)
+
+    else:
+        # If DS is provided in CLI instead of json, use it as a list (default)
+        dataset = dataset_arg
+
+    return dataset
+
+
+@app.command()
+def run_from_command(
+    dataset: List[str] = typer.Argument(
+        ...,
+        help="Input datasets (Rucio DID) or path to JSON file containing datasets in a dict.",
+    ),
+    filter_branch: str = typer.Option(
+        "", "--filter-branch", help="Only display branches containing this string."
+    ),
+):
+    """
+    Calls the get_structure function and sends results to stdout.
+    To run on command line: servicex-get-structure -dataset --filter-branch
+    """
+    ds_format = make_dataset_list(dataset)
+    result = get_structure(ds_format, filter_branch=filter_branch, do_print=False)
+
+    print(result)
+
+
+if __name__ == "__main__":
+    app()