feat: update examples/pipeline_2 to align with package better

examples/pipeline_2/conf.yaml

@@ -17,7 +17,7 @@ pipeline_variables:
   working_dir: examples/pipeline_2
   project_root: examples/pipeline_2
   data_dir: examples/pipeline_2/data
-  output_dir: examples/pipeline_2/outputs
+  output_dir: examples/pipeline_2/runs
   log_dir: examples/pipeline_2/logs
   overwrite: True
   metadata:
@@ -48,11 +48,7 @@ stage_configuration:
       - "age"
       - "dob"
       - "address"
-    output_location: examples/pipeline_2/data/orders_cleaned.csv
-    input_location: examples/pipeline_2/data/orders.csv
   1_derive_vars:
-    input_location: examples/pipeline_2/data/orders_cleaned.csv
-    output_location: examples/pipeline_2/data/orders_prepped.csv
     delivery_times:
       north: 14
       south: 4
@@ -82,8 +78,6 @@ stage_configuration:
     total_production_cost: "Total_production_cost"
     order_profit: "Order_profit"
   2_reporting:
-    input_location: examples/pipeline_2/data/orders_prepped.csv
-    report_location: examples/pipeline_2/outputs/order_analysis.md
     region: "Region"
     order_profit: "Order_profit"
     quantity: "Quantity"

examples/pipeline_2/scripts/0_clean_data.py

@@ -31,21 +31,39 @@ def standardise_columns(df):
     return df
 
 
-def main(context=None):
+def main(context):
+    # Load stage configuration
     config = context.get_stage_config("0_clean_data")
-    print(config)
 
-    orders = pd.read_csv(config["input_location"])
+    # Calculate location for run outputs
+    output_root = context.resolve_output_root()
+    full_output_location = output_root / "orders_cleaned.csv"
 
+    # Calculate location of data input
+    data_dir = context.get_data_dir()
+    orders = pd.read_csv(data_dir / "orders.csv")
+
+    # Source variable lists from stage configuration
     expected_variables = config["expected_variables"]
     identifiable_cols = config["identifiable_cols"]
 
+    # Run functions required for this stage
     check_variables(orders, expected_variables)
     print(orders.dtypes)
     orders = remove_identifiable(orders, identifiable_cols)
     orders = standardise_columns(orders)
-    orders.to_csv(config["output_location"], index=False)
+
+    # Save cleaned data to run specific output location defined earlier in function
+    orders.to_csv(full_output_location, index=False)
+
+    return {
+        "output_location": str(full_output_location),
+        "expected_variables": expected_variables,
+        "identifiable_cols": identifiable_cols,
+        "record_count": len(orders),
+        "columns": list(orders.columns),
+    }
 
 
 if __name__ == "__main__":
-    main()
+    main(context=None)

examples/pipeline_2/scripts/1_derive_vars.py

@@ -57,12 +57,21 @@ def profit_per_order(df):
     return df
 
 
-def main(context=None):
+def main(context):
+    # Get stage configuration for this stage
     config = context.get_stage_config("1_derive_vars")
 
-    df = pd.read_csv(config["input_location"])
+    # Source data path for previous stage results which are required
+    # for this stage and read those in
+    data_path = context.resolve_given_path(
+        "0_clean_data", "output_location", "orders_cleaned.csv", context.get_data_dir()
+    )
+    df = pd.read_csv(data_path)
+
+    # Source variable list from stage configuration for this stage
     delivery_times = config["delivery_times"]
 
+    # Run relevant functions for this stage
     df = correct_date_time(df)
     df = estimate_delivery(df, delivery_times)
     df = total_cost(df)
@@ -71,7 +80,16 @@ def main(context=None):
     df = postage_cost(df)
     df = production_cost(df)
     df = profit_per_order(df)
-    df.to_csv(config["output_location"], index=False)
+
+    # Save results to output location calculated based on run_directory
+    output_root = context.resolve_output_root()
+    df.to_csv(output_root / "orders_prepped.csv", index=False)
+
+    return {
+        "output_location": str(output_root / "orders_prepped.csv"),
+        "record_count": len(df),
+        "columns": list(df.columns),
+    }
 
 
 if __name__ == "__main__":

examples/pipeline_2/scripts/2_reporting.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 import pandas as pd
 
 
@@ -116,20 +114,27 @@ def curate_report(report, values):
     )
 
 
-def write_report(report):
-    report_file = Path("examples/pipeline_2/outputs/order_analysis.md")
+def write_report(report, output_path):
+    report_file = output_path
 
     report_file.write_text("\n".join(report), encoding="utf-8")
 
 
-def main():
-    orders = pd.read_csv("examples/pipeline_2/data/orders_prepped.csv")
+def main(context):
+    # Read in processed data from stage "1_derive_vars"
+    data_loc = context.resolve_given_path(
+        "1_derive_vars", "output_location", "orders_prepped.csv", context.get_data_dir()
+    )
+    orders = pd.read_csv(data_loc)
 
+    # Set empty list and dictionary to store report and values
     report = []
     values = {}
 
+    # Set the number format for the report to 2 decimal places
     num_format = "{:.2f}"
 
+    # Run functions required for the stage
     per_region_profits(orders, values, num_format)
     per_region_quantity(orders, values)
     orders_per_day(orders, values)
@@ -138,7 +143,16 @@ def main():
     total_summaries(orders, values, num_format)
     profit_per_product(orders, values, num_format)
     curate_report(report, values)
-    write_report(report)
+
+    # Calculate the output location for the final report and write to it
+    output_path = context.resolve_output_root() / "order_analysis.md"
+    write_report(report, output_path)
+
+    return {
+        "report_location": str(output_path),
+        "total_orders": int(values["total_count"]),
+        "total_profit": float(values["total_profit"]),
+    }
 
 
 if __name__ == "__main__":