| from typing import Dict |
| |
| import my_functions |
| from pandas import DataFrame |
| |
| from hamilton import base, driver, lifecycle |
| |
| driver = ( |
| driver.Builder() |
| .with_config({}) |
| .with_modules(my_functions) |
| .with_adapters( |
| # this is a strict type checker for the input and output of each function. |
| lifecycle.FunctionInputOutputTypeChecker(), |
| # this will make execute return a pandas dataframe as a result |
| base.PandasDataFrameResult(), |
| ) |
| .build() |
| ) |
| |
| |
| class TaxCalculator: |
| """ |
| Simple class to wrap Hamilton Driver |
| """ |
| |
| @staticmethod |
| def calculate( |
| input: DataFrame, tax_rates: Dict[str, float], tax_credits: Dict[str, float] |
| ) -> DataFrame: |
| return driver.execute( |
| inputs={"input": input, "tax_rates": tax_rates, "tax_credits": tax_credits}, |
| final_vars=["final_tax_dataframe"], |
| ) |
| |
| @staticmethod |
| def visualize(output_path="./my_full_dag.png"): |
| # To visualize do `pip install "sf-hamilton[visualization]"` if you want these to work |
| driver.display_all_functions(output_path) |
