Sequential Tasks in Flow with Mapping of Task Output in Imperative API

[taylorfturner]

Context

See #4028 for context. This discussion is an extension of that: trying to port over the functional API Flow into the imperative API paradigm.

Issue

AttributeError: <Task: ProcessData> does not have data_frame as an attribute

Expected Behavior

I expect the same behavior in the imperative API flow that the functional API flow example in the code snippet and the block directly below (a copy+p of from the code snippet block).

data_processing = ProcessData()
reader = DataReader()
data_list=['data_one', 'data_two', 'data_three']

with Flow('imperative flow with mapping') as functional_flow:

    data_type = Parameter('data_type', default='pandas')
    functional_flow.add_task(data_type)

    data = reader.map(
        data_type=unmapped(data_type), 
        data_list=data_list
    )
    data_processing.map(data_frame=data)

functional_flow.run()

Code Snippet

The functional API example works. The imperative API example does not work (yet 🤞).

from prefect import Flow, Parameter, unmapped, Task
import pandas as pd
import numpy as np


class DataReader(Task):
    def _read_pandas(self):
        df = pd.DataFrame(np.random.randn(100, 4), columns=list("ABCD"))
        return df

    def _read_dask(self):
        raise NotImplementedError

    def run(self, data_type, data_list):
        if data_type == "dask":
            return self._read_dask()
        if data_type == "pandas":
            return self._read_pandas()


class ProcessData(Task):
    def __init__(self):
        super().__init__()

    def run(self, data_frame):
        return data_frame.apply(lambda x: (x + 1).all())



# this doesn't work
# how do I `.map()` through the imperative flow API?
class ImperativeFlow(Flow):
    def __init__(self, name):
        super().__init__(name)

    def get_flow(self, data_list=['data_one', 'data_two', 'data_three']):

        reader = DataReader()
        process_data = ProcessData()

        data_type = Parameter('data_type', default='pandas')
        self.add_task(data_type)
        self.set_dependencies(
            task=process_data(
                mapped=True,
                task_args={'data_frame': reader}
                ),
            upstream_tasks=reader(
                mapped=True,
                task_args={'data_type': unmapped(data_type), 'data_list': data_list})
        )
        return self

    def run(self):
        self.run()




# we got this work. I'm tring to replicate the same final
# graph using the imperative API in a Class
data_processing = ProcessData()
reader = DataReader()
data_list=['data_one', 'data_two', 'data_three']

with Flow('imperative flow with mapping') as functional_flow:

    data_type = Parameter('data_type', default='pandas')
    functional_flow.add_task(data_type)

    data = reader.map(
        data_type=unmapped(data_type), 
        data_list=data_list
    )
    data_processing.map(data_frame=data)

functional_flow.run()