feat: pick return values

Makefile

@@ -64,10 +64,13 @@ generate-test-outputs:
 	@rm -rf tests/outputs_basic_with_hooks
 	@uv run efemel process "**/*.py" --cwd tests/inputs/basic --out tests/outputs/with_hooks --hooks tests/hooks/before_after/output_filename.py
 
-	@echo "Generating test outputs with hooks directory..."
 	@rm -rf tests/outputs/with_hooks_dir
 	@uv run efemel process "**/*.py" --cwd tests/inputs/basic --out tests/outputs/with_hooks_dir --hooks tests/hooks/multiple
 
+	@rm -rf tests/outputs/with_hooks_dir
+	@uv run efemel process "**/*.py" --cwd tests/inputs/process_data --out tests/outputs/process_data --pick user_data
+
+
 # Clean build artifacts and cache files
 clean:
 	rm -rf build/

efemel/cli.py

@@ -4,7 +4,8 @@
 
 import click
 
-from efemel.hooks.output_filename import ensure_output_path, flatten_output_path
+from efemel.hooks import output_filename
+from efemel.hooks import process_data as process_data_hooks
 from efemel.hooks_manager import HooksManager
 from efemel.process import process_py_file
 from efemel.readers.local import LocalReader
@@ -49,7 +50,13 @@ def info():
   type=click.Path(exists=True, readable=True, resolve_path=True),
   help="Path to a Python file or directory containing user-defined hooks.",
 )
-def process(file_pattern, out, flatten, cwd, env, workers, hooks):
+@click.option(
+  "--pick",
+  "-p",
+  multiple=True,
+  help="Pick specific dictionary keys to extract (can be used multiple times)",
+)
+def process(file_pattern, out, flatten, cwd, env, workers, hooks, pick):
   """Process Python files and extract public dictionary variables to JSON.
 
   FILE_PATTERN: Glob pattern to match Python files (e.g., "**/*.py")
@@ -63,7 +70,10 @@ def process(file_pattern, out, flatten, cwd, env, workers, hooks):
 
   if flatten:
     # Add the flatten_output_path hook to the hooks manager
-    hooks_manager.add("output_filename", flatten_output_path)
+    hooks_manager.add("output_filename", output_filename.flatten_output_path)
+
+  if pick:
+    hooks_manager.add("process_data", process_data_hooks.pick_data(pick))
 
   # Load user-defined hooks if a path is specified
   if hooks:
@@ -76,7 +86,7 @@ def process(file_pattern, out, flatten, cwd, env, workers, hooks):
     else:
       click.echo(f"Warning: Hooks path '{hooks}' is neither a file nor a directory")
 
-  hooks_manager.add("output_filename", ensure_output_path)
+  hooks_manager.add("output_filename", output_filename.ensure_output_path)
 
   # Collect all files to process
   files_to_process = list(reader.read(file_pattern))
@@ -90,7 +100,17 @@ def process_single_file(file_path: Path, cwd: Path):  # Added type hint for clar
     try:
       # Always create output file, even if no dictionaries found
       public_dicts = process_py_file(cwd / file_path, env) or {}
-      transformed_data = transformer.transform(public_dicts)
+
+      (processed_data,) = hooks_manager.call(
+        "process_data",
+        {
+          "data": public_dicts,
+          "env": env,
+        },
+        return_params=["data"],
+      )
+
+      transformed_data = transformer.transform(processed_data)
 
       # Original proposed output filename (as a Path object for consistency)
       proposed_output_path = file_path.with_suffix(transformer.suffix)

efemel/hooks/process_data.py

@@ -0,0 +1,9 @@
+def pick_data(keys):
+  """Pick specific keys from the processed python file."""
+
+  def _pick(context):
+    if not keys:
+      return
+    context["data"] = {key: value for key, value in context["data"].items() if key in keys}
+
+  return _pick

tests/inputs/process_data/data.py

@@ -0,0 +1,16 @@
+config = {
+  "name": "test_app",
+  "version": "1.0.0",
+  "debug": True,
+  "features": ["auth", "api", "ui"],
+}
+
+settings = {
+  "database": {"host": "localhost", "port": 5432, "name": "app_db"},
+  "cache": {"type": "redis", "ttl": 3600},
+}
+
+user_data = {
+  "admin": {"name": "Admin User", "role": "admin"},
+  "guest": {"name": "Guest User", "role": "guest"},
+}

tests/outputs/process_data/data.json

@@ -0,0 +1,12 @@
+{
+  "user_data": {
+    "admin": {
+      "name": "Admin User",
+      "role": "admin"
+    },
+    "guest": {
+      "name": "Guest User",
+      "role": "guest"
+    }
+  }
+}

tests/test_cli.py

@@ -42,12 +42,18 @@ def get_test_scenarios():
       "hooks": ["before_after/output_filename.py"],
     },
     {
-      "name": "hooks-dir",
+      "name": "hooks dir",
       "inputs_dir": test_dir / "inputs/basic",
       "outputs_dir": test_dir / "outputs/with_hooks_dir",
       "process_args": ["--hooks", "hooks/multiple"],
       "hooks": ["multiple/output_filename.py"],
     },
+    {
+      "name": "process data - pick",
+      "inputs_dir": test_dir / "inputs/process_data",
+      "outputs_dir": test_dir / "outputs/process_data",
+      "process_args": ["--pick", "user_data"],
+    },
   ]