From de9d3a035e8f26a5e9e1f22bab5d8c26cfe97db6 Mon Sep 17 00:00:00 2001
From: Orion Yeung <11580988+orionyeung001@users.noreply.github.com>
Date: Sat, 25 Nov 2023 21:52:22 -0600
Subject: [PATCH] add --field-prefix via click.option and update doc

added option proposed in scikit-hep/histoprint #121

implement behavior of --field-prefix option

uses click callback on the field option to handle this.

updated docs to include help text for --field-prefix option

remove type annotations

One of these type annotations was actually incorrect, but I noticed
that they weren't used in this file. Opted to remove instead of fix
---
 README.rst        |  6 ++++++
 histoprint/cli.py | 17 ++++++++++++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/README.rst b/README.rst
index 8334a79..93ab3f8 100644
--- a/README.rst
+++ b/README.rst
@@ -179,6 +179,12 @@ It can read in files or take data directly from STDIN::
                                       'tree/branch[:,2]' to histogram the 3rd
                                       elements of a vector-like branch.
 
+      --field-prefix TEXT             String to prepend to each argument passed to
+                                      --field option following --field-prefix E.g.
+                                      'histoprint -f Muon_eta --field-prefix Tau_
+                                      -f eta -f phi' ... would refer to fields
+                                      'Muon_eta', 'Tau_eta' and 'Tau_phi''
+
       -C, --cut TEXT                  Filter the data to be plotted by a cut
                                       condition. For ROOT files, variables must be
                                       referenced by their branch name within the
diff --git a/histoprint/cli.py b/histoprint/cli.py
index a5e7e39..e7ec3c8 100644
--- a/histoprint/cli.py
+++ b/histoprint/cli.py
@@ -9,6 +9,11 @@
 import histoprint as hp
 import histoprint.formatter as formatter
 
+def fieldprefixer_callback(ctx, param, fields):
+    if 'field_prefix' in ctx.params:
+        fields = map(lambda field: ctx.params['field_prefix'] + field, fields)
+    return tuple(fields)
+
 
 @click.command()
 @click.argument("infile", type=click.Path(exists=True, dir_okay=False, allow_dash=True))
@@ -74,6 +79,16 @@
     "single TH1, or one or more paths to TTree branches. Also supports slicing "
     "of array-like branches, e.g. use 'tree/branch[:,2]' to histogram the 3rd "
     "elements of a vector-like branch.",
+    callback=fieldprefixer_callback
+)
+@click.option(
+    "--field-prefix",
+    "field_prefix",
+    type=str,
+    multiple=False,
+    help="String to prepend to each argument passed to --field option following --field-prefix"
+    "\nE.g. `histoprint -f Muon_eta --field-prefix Tau_ -f eta -f phi` ... would refer to fields"
+    "\n\t`Muon_eta`, `Tau_eta` and `Tau_phi`'",
 )
 @click.option(
     "-C",
@@ -109,6 +124,7 @@ def histoprint(infile, **kwargs):
 
     INFILE can be '-', in which case the data is read from STDIN.
     """
+    del kwargs["field_prefix"]
 
     # Read file into buffer for use by implementations
     try:
@@ -148,7 +164,6 @@ def histoprint(infile, **kwargs):
     click.echo("Could not interpret file format.", err=True)
     exit(1)
 
-
 def _bin_edges(kwargs, data):
     """Get the desired bin edges."""
     bins = kwargs.pop("bins", "10")