blog_20140805_5.py

###### Meta class ######

class DependencyMetaTask(luigi.Task):
    # METHODS FOR AUTOMATING DEPENDENCY MANAGEMENT
    def requires(self):
        upstream_tasks = []
        for param_val in self.param_args:
            if type(param_val) is dict:
                if 'upstream' in param_val:
                    upstream_tasks.append(param_val['upstream']['task'])
        return upstream_tasks
 
    def get_input(self, input_name):
        param = self.param_kwargs[input_name]
        if type(param) is dict and 'upstream' in param:
            return param['upstream']['task'].output()[param['upstream']['port']]
        else: 
            return param


###### Normal classes ######

class TaskA(DependencyMetaTask):
    # INPUT TARGETS
    in1_target = luigi.Parameter()
    param_a1 = luigi.Parameter()

   # DEFINE OUTPUTS
    def output(self):
        return { 'out1' : 
            luigi.LocalTarget(
                self.get_input('in1_target').path + '.out1'),
                 'out2' : 
            luigi.LocalTarget(
                self.get_input('in1_target').path + '.out2') } }

    # WHAT THE TASK DOES
    def run(self):
        with open(self.get_input('in1_target').path) as infile:
            for line in infile:
                do_something(line)


class TaskB()
    # INPUT TARGETS
    in1_target = luigi.Parameter()
    in2_target = luigi.Parameter()
    param_b1 = luigi.Parameter()
    param_b2 = luigi.Parameter()

    def run(self):
        # Do something with both in1 and in2
        ....

##### THE ACTUAL WORKFLOW / DEPENDENCY GRAPH DEFINITION #####

class MyWorkFlow(luigi.Task):
    # We only need to duplicate all parameters 
    # once, which is here in the workflow task
    param_a1 = luigi.Parameter()
    param_b1 = luigi.Parameter()
    param_b2 = luigi.Parameter()
 
    # Here the whole workflow definition resides:
    def requres(self):
        task_a = TaskA( 
            param_a1 = self.param_a1 
        )

        task_b = TaskB( 
            param_b1 = self.param_b1,
            param_b2 = self.param_b2,
            # Here below, we connect the output out1 from TaskA
            # to in1_target of TaskB ...
            in1_target = 
                { 'upstream' : { 'task' : task_a,
                                 'port' : 'out1' } }
            # ... and again, out2 of TaskA, to in2_target of
            # TaskB, using our special syntax.
            in2_target = 
                { 'upstream' : { 'task' : task_a,
                                 'port' : 'out2' } }
        )