Make context less nebulous

@kaxil

Description

Currently context is a dictionary. The dictionary is a bit nebulous especially for new users as it isn't clear what is present inside it and the autocomplete functionality of arbitrary dictionaries leaves a lot to be desired.

Discoverability and type safety could be greatly improved if context was refactored into a dataclass (preferred) or typedDict(less ideal but probably easier).

Use case / motivation

Avoid having to know the keys present in context

def get_current_context() -> Dict:
  pass
  

context = get_current_context()
ti = context["ti"]

would be instead

@dataclass
class Context:
  # would be the place to see what values are in context and optionally include comments
  ti: TaskInstance # this is the task Instance
  ds: DateTime # This is the execution date
  # etc...

def get_current_context() -> Context:
  pass
  
context = get_current_context()
#`context.` would have strong autocomplete here
ti = context.ti

Context Dataclass MVP V2 From later in the thread that prevents a breaking change & warns on unsafe updates to context

@dataclass
class Demo:  # context replacement
    id: str
    value_dc: int
    user_defined: Dict[str, Any] = field(default_factory=dict)

    def __getitem__(self, item):
        if item in self.__dict__.keys():
            logging.warning(msg=f"dictionary interface getitem on context is deprecated; update to use the dataclass interface for standard fields like `{item}`")
            return self.__dict__[item]
        elif item in self.user_defined:
            logging.warning(msg=f"dictionary interface getitem on context is deprecated; update to use context.user_defined for custom fields like `{item}`")
            return self.user_defined[item]
        else:
            raise KeyError

    def __setitem__(self, key: str, value):
        if key in self.__dict__.keys():
            msg = f"""dictionary interface setitem for standard fields is deprecated; update to use the dataclass interface for standard fields like `{key}`
            note: changing standard context fields is not supported and may have undefined behavior. If this is meant to be a custom field use context.user_defined instead"""
            logging.warning(msg=msg)
            self.__dict__[key] = value
        else:
            logging.warning(
                msg=f"dictionary interface setitem on context is deprecated; update to use context.user_defined for custom fields like `{key}`")
            self.user_defined[key] = value

    def keys(self):
        # added as an example to show how far we could go to have a non-breaking change for 2.1
        logging.warning(msg=f"dictionary interface keys is deprecated; update this to use the dataclass interface")
        temp = self.__dict__
        temp.update(self.user_defined)
        return temp


d = Demo(id="long_id", value_dc=1337)
print(d["id"])
d["new"] = 3
print(d["new"])
print(d.keys())
d["id"] = "warn"

returns

WARNING:root:dictionary interface getitem on context is deprecated; update to use the dataclass interface for standard fields like `id`
WARNING:root:dictionary interface setitem on context is deprecated; update to use context.user_defined for custom fields like `new`
WARNING:root:dictionary interface getitem on context is deprecated; update to use context.user_defined for custom fields like `new`
WARNING:root:dictionary interface keys is deprecated; update this to use the dataclass interface
WARNING:root:dictionary interface setitem for standard fields is deprecated; update to use the dataclass interface for standard fields like `id`
            note: changing standard context fields is not supported and may have undefined behavior. If this is meant to be a custom field use context.user_defined instead
long_id
3
{'id': 'long_id', 'value_dc': 1337, 'user_defined': {'new': 3}, 'new': 3}

My end goal is to be able to do the following when writing callables that can have context passed to them

def generate_is_latest_callable(tasks_if_latest: List[str] , tasks_if_not_latest: List[str]) -> Callable:
  def result(context: Context) -> List[str]:
    context. #and be able to get strong autocomplete and typing while in an ide here
    # because currently all I can do is context: Dict[str, Any] which isn't very helpful
    if context.something:
      return tasks_if_latest
    else:
      return tasks_if_not_latest
  return result

t_branch = PythonBranchingOperator(
  task_id="branch",
  python_callable=generate_is_latest_callable(["yes"], ["no"]),
  provide_context=True,
  dag=dag,
)

Are you willing to submit a PR?

This one touches too much of airflow's internals for me to try and tackle.

Related Issues

Not to my knowledge
Slack Conversation: https://apache-airflow.slack.com/archives/C0146STM600/p1614103441043600?thread_ts=1614099549.042800&cid=C0146STM600

@kaxil

I think context as a ductionary is pretty sound. Dataclass is not pythonic at all and makes it difficult to evolve. I think it's very pythonic to have dictionaries in Python and it brings a number of benefits, with little drawback

Dataclass is not pythonic at all

As a counter point Dacite is a python library that shows how dataclasses can be used for something like this. Their config dataclass is a similar object to our context and is defined as follows.

https://github.com/konradhalas/dacite/blob/d2206b2e4711859da0ea5862c395940f33693e80/dacite/config.py#L5-L12

@dataclass
class Config:
    type_hooks: Dict[Type, Callable[[Any], Any]] = field(default_factory=dict)
    cast: List[Type] = field(default_factory=list)
    forward_references: Optional[Dict[str, Any]] = None
    check_types: bool = True
    strict: bool = False
    strict_unions_match: bool = False

which is then consumed in their function from_dict
https://github.com/konradhalas/dacite/blob/d2206b2e4711859da0ea5862c395940f33693e80/dacite/core.py#L34

def from_dict(data_class: Type[T], data: Data, config: Optional[Config] = None) -> T:
    """Create a data class instance from a dictionary.
    :param data_class: a data class type
    :param data: a dictionary of a input data
    :param config: a configuration of the creation process
    :return: an instance of a data class

As a user it feels so much better to be able to see the config in the function definition and then follow the type definition straight to the config objects strict definition (fields, types, optionally comments).

Alternatively when a dictionary is used its harder for an end user to reason about what keys are present?, what are the values?, what are the types of the values?, where do I find this information?. (Perhaps its my own ignorance but it isn't obvious to me today where I can look up this info for airflow's context dictionary and I've tried).

As a developer tools like mypy, pydantic & various ides (pycharm, vscode) also handle dataclasses quite well which helps prevent frustrating runtime errors.

makes it difficult to evolve

Not sure about this one.

To add new fields to data classes you simply add a new field to the dataclass definition with a default value.
To remove old fields you can use an ide's find references to remove all the old references and then delete the field from the dataclass.
- Optionally, use a static analyzer like mypy before commit/merge to ensure you've gotten rid of all the old references

Perhaps I'm still a rookie but for dictionaries it feels much harder to be sure you've gotten all the old key references out of the code, before you can safely remove it.

I think it's very pythonic to have dictionaries in Python and it brings a number of benefits, with little drawback

I've called out a couple advantages I think dataclasses have over dictionaries. But I'm likely missing the dictionary advantages. Is Airflow using any of these advantages today?

We could also add a TypedDict -- at least IDEs will have good auto-completion

@kaxil

We could also add a TypedDict -- at least IDEs will have good auto-completion

Yep. @kaxil is right - TypedDict is far more pythonic (Hints are the way how types are usually added in Python). The thing with using fixed Dataclasses. We should basically never remove/rename anything in the context (maybe we should use some deprecation mechanism).

The context is used by pretty much every single custom operator out there, and the change you propose will immediately break all of them.. If anything else - this is a single reason why we should not do it.

But the idea of improving developer's productivity is sound and since we have TypedDict which is backwards compatible and gives most benefits you mentioned, I think it is far better solution. Would youl like to pick this up @r-richmond ?

@r-richmond

TypedDict is far more pythonic (Hints are the way how types are usually added in Python)

Hints are how dataclasses uses types as well. They aren't enforced at runtime, unless you use something like pydantic. (side note Dataclass PEP-557 had support from Raymond & Guido which as far as I can tell gives it a lot of points in the pythonic category. Also Raymond gave a fantastic talk on Dataclasses that I'd recommend if you have time).

We should basically never remove/rename anything in the context (maybe we should use some deprecation mechanism).

Yep makes sense to me. Although, this isn't a point for or against dataclasses, typedDicts, or dicts.

The context is used by pretty much every single custom operator out there, and the change you propose will immediately break all of them.. If anything else - this is a single reason why we should not do it.

While switching to dataclasses would be harder than TypedDicts it does not have to be a breaking change. You can convert dictionaries to dataclasses and vice versa pretty seamlessly. Although, implementing Dataclasses would certainly be more work to get right in the initial pr.

Example showing how you can switch between dataclasses & dictionaries

from dataclasses import dataclass, asdict
from dacite import from_dict


@dataclass
class Demo:
  id: str
  value: int


d = Demo("a", 1)

examples = [d, asdict(d), from_dict(data_class=Demo, data=asdict(d))]
for example in examples:
  print(example, type(example))

prints

Demo(id='a', value=1) <class '__main__.Demo'>
{'id': 'a', 'value': 1} <class 'dict'>
Demo(id='a', value=1) <class '__main__.Demo'>

Personally I'm of the opinion that working with the dataclass api is superior to working with the dictionary api due to the reasons outlined above but thats just like my opinion man.

But the idea of improving developer's productivity is sound and since we have TypedDict which is backwards compatible and gives most benefits you mentioned

Either option could be backwards compatible but I agree that the TypedDict change is more straight forward and is an alternative, hence the tiny blurb about it in the opener.

As you mentioned TypedDict should give most of the benefits, although to my knowledge find references does not work as well dictionary keys which is a bummer. Anyways the decision on which to use is above my Airflow status so I'll let the maintainers decide. Most importantly I think either a TypedDict or Dataclass implementation would be a big improvement over the current situation.

Would youl like to pick this up @r-richmond ?

Actually, the reason I submitted this issue was because me and a coworker were frustrated with the "magic context dictionary". TLDR I'm not actually sure what all the keys, values, & types are in context so I wouldn't even know how to start this one.

My end goal is to be able to do the following when writing callable that can have context passed to them

def generate_is_latest_callable(tasks_if_latest: List[str] , tasks_if_not_latest: List[str]) -> Callable:
  def result(context: Context) -> List[str]:
    context. #and be able to get strong autocomplete and typing while in an ide here
    # because currently all I can do is context: Dict[str, Any] which isn't very helpful
    if context.something:
      return tasks_if_latest
    else:
      return tasks_if_not_latest
  return result

t_branch = PythonBranchingOperator(
  task_id="branch",
  python_callable=generate_is_latest_callable(["yes"], ["no"]),
  provide_context=True,
  dag=dag,
)

Note: moving this goal to the opening comment

Oh yeah, I will add docs (in few mins / hours) to make the context dictionary a bit less magical, it contains the following items:

airflow/airflow/models/taskinstance.py

Lines 1688 to 1728 in 352b970

    
           return { 
        
               'conf': conf, 
        
               'dag': task.dag, 
        
               'dag_run': dag_run, 
        
               'ds': ds, 
        
               'ds_nodash': ds_nodash, 
        
               'execution_date': pendulum.instance(self.execution_date), 
        
               'inlets': task.inlets, 
        
               'macros': macros, 
        
               'next_ds': next_ds, 
        
               'next_ds_nodash': next_ds_nodash, 
        
               'next_execution_date': next_execution_date, 
        
               'outlets': task.outlets, 
        
               'params': params, 
        
               'prev_ds': prev_ds, 
        
               'prev_ds_nodash': prev_ds_nodash, 
        
               'prev_execution_date': prev_execution_date, 
        
               'prev_execution_date_success': lazy_object_proxy.Proxy( 
        
                   lambda: self.get_previous_execution_date(state=State.SUCCESS) 
        
               ), 
        
               'prev_start_date_success': lazy_object_proxy.Proxy( 
        
                   lambda: self.get_previous_start_date(state=State.SUCCESS) 
        
               ), 
        
               'run_id': run_id, 
        
               'task': task, 
        
               'task_instance': self, 
        
               'task_instance_key_str': ti_key_str, 
        
               'test_mode': self.test_mode, 
        
               'ti': self, 
        
               'tomorrow_ds': tomorrow_ds, 
        
               'tomorrow_ds_nodash': tomorrow_ds_nodash, 
        
               'ts': ts, 
        
               'ts_nodash': ts_nodash, 
        
               'ts_nodash_with_tz': ts_nodash_with_tz, 
        
               'var': { 
        
                   'json': VariableJsonAccessor(), 
        
                   'value': VariableAccessor(), 
        
               }, 
        
               'yesterday_ds': yesterday_ds, 
        
               'yesterday_ds_nodash': yesterday_ds_nodash, 
        
           }

Ah interesting kaxil thanks to your link I can see where a TypedDict / Dataclass implementation would go

airflow/airflow/models/taskinstance.py

Line 86 in 352b970

Context = Dict[str, Any]

Since the TypedDict implementation is so straightforward given that link I wonder if its better to consider Dict > TypeDict > DataClass as an evolution path.

While switching to dataclasses would be harder than TypedDicts it does not have to be a breaking change. You can convert dictionaries to dataclasses and vice versa pretty seamlessly. Although, implementing Dataclasses would certainly be more work to get right in the initial pr.

Let me explain why this would be a breaking change. if you have an operator that already uses context. Example from task decorator case (but this is the same for any custom operator):

from airflow.operators.python import task, get_current_context

@task
def my_task():
    context = get_current_context()
    ti = context["ti"]

Surely you can map from Dataclass to Dict in "get_current_context()" method (otherwise you'd have to also correct the context["ti"] into context.ti. Thus it would be a breaking change.

But if you do that (i.e map dataclass to dict), then you loose the type hinting. This is precisely the place where typehinting is needed (in the user code). And I think Typed Dict (if used everywhere including return value from get_current_context() gives you what you need from the user point of view and there is no need to implement the > Dataclass step - as you have everything you need by implementing TypeDict.

Or maybe I am missing something?

Personally I'm of the opinion that working with the dataclass api is superior to working with the dictionary api due to the reasons outlined above but thats just like my opinion man.

In a number of cases yes. But in others flexibility of the dict and fact that you can add any value there, trumps.

Also one more point - since we are in Python world - users could do whatever they want in the Python code of their custom operators. Currently 'context' is part of the public API of Airflow that users rely on being dictionary.

For example I can very easily imagine someone decides to add a new field to context and pass it further.

@task
def my_task():
    context = get_current_context()
    ti = context["ti"]
    context['my_key'] = 'my_value'
    call_another_function(context)

While we do not encourage this behavior in any way, it is tempting to use the context this way and as a user you could do that and you could rely on context being dictionary able to hold any key you want. Dataclass breaks this assumption.

Add cross-reference for Context Dictionary (#14428)

Navigation Menu

Repository navigation

Make context less nebulous #14396

46 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	return {
	'conf': conf,
	'dag': task.dag,
	'dag_run': dag_run,
	'ds': ds,
	'ds_nodash': ds_nodash,
	'execution_date': pendulum.instance(self.execution_date),
	'inlets': task.inlets,
	'macros': macros,
	'next_ds': next_ds,
	'next_ds_nodash': next_ds_nodash,
	'next_execution_date': next_execution_date,
	'outlets': task.outlets,
	'params': params,
	'prev_ds': prev_ds,
	'prev_ds_nodash': prev_ds_nodash,
	'prev_execution_date': prev_execution_date,
	'prev_execution_date_success': lazy_object_proxy.Proxy(
	lambda: self.get_previous_execution_date(state=State.SUCCESS)
	),
	'prev_start_date_success': lazy_object_proxy.Proxy(
	lambda: self.get_previous_start_date(state=State.SUCCESS)
	),
	'run_id': run_id,
	'task': task,
	'task_instance': self,
	'task_instance_key_str': ti_key_str,
	'test_mode': self.test_mode,
	'ti': self,
	'tomorrow_ds': tomorrow_ds,
	'tomorrow_ds_nodash': tomorrow_ds_nodash,
	'ts': ts,
	'ts_nodash': ts_nodash,
	'ts_nodash_with_tz': ts_nodash_with_tz,
	'var': {
	'json': VariableJsonAccessor(),
	'value': VariableAccessor(),
	},
	'yesterday_ds': yesterday_ds,
	'yesterday_ds_nodash': yesterday_ds_nodash,
	}

Make context less nebulous #14396

Description

Activity

potiuk commented on Feb 23, 2021

r-richmond commented on Feb 24, 2021

kaxil commented on Feb 24, 2021

potiuk commented on Feb 24, 2021

r-richmond commented on Feb 24, 2021

kaxil commented on Feb 24, 2021

r-richmond commented on Feb 24, 2021

potiuk commented on Feb 24, 2021

potiuk commented on Feb 24, 2021

46 remaining items

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions