Context class handles deprecation

[uranusjr]

Fix #19716, close #19734.

Fortunately (?) we don’t really set up Mypy with enough information to catch typing errors this introduces to def execute(self, context: Dict):. I still intend to fix those eventually, but this helps producing a cleaner diff to run against CI.

Now that we have a class (and a separate module), I also want to move more logic from the giant TaskInstance.get_template_context() (currently ~250 lines).

[uranusjr]

I wonder if this would help #14396. We can make Context a TypedDict, but I really don’t like subclassing a built-in class (it’s usually a cause for weird bugs).

[potiuk]

@uranusjr -> FYI MyPy is disabled temporarily, in case you have not noticed. But I think I will add it back in non-failure mode, just to incrementally fix the remaining problems.

[uranusjr]

Bad news, you can’t override anything in TypedDict. So we need an alternative way to handle #14396…

[potiuk]

Bad news, you can’t override anything in TypedDict. So we need an alternative way to handle #14396…

How about we use dynamic attributes with set/get_attr and stub files from PEP 484 ? https://www.python.org/dev/peps/pep-0484/#stub-files - I think major IDEs already nicely support them. This way we would not have to unnecesarily complicate the implementation.

That very nicely fits my "requirements" for less-nebulous context #14396 (comment)

[uranusjr]

If we use stub files, we can continue using a dict-like class (i.e. don’t need dynamic setattr/getattr) and just “pretend” the class is actually a TypedDict in .pyi. That should work.

[ferruzzi]

As mentioned in #14396, I have a chunk of that one already done in my working branch. I wanted to try to polish it a bit more before sending in a PR; it currently passes local CI and supports type hinting, but not auto-completing.

As you pointed out, TypedDict won't let you override anything. I had hoped to use @Property tags to allow auto-completing, but no dice there. I'll look into stub files, I'm not familiar with using those yet.

[uranusjr]

If I manage to get the tests pass for the Context class change (trying to do that right now), I’ll add a stub file tomorrow as an example how this would work that you can test out.

[ferruzzi]

Did some reading and checked out stubgen which looked like it basically just replicated the work I did already... which was both neat and depressing. :P I don't want to clutter your thread here, but please feel free to reach out to me on Slack, perhaps?

[uranusjr]

I pushed a commit with a context.pyi file that defines Context as a TypedDict. This makes VSCore + Python extension trigger autocompletion on context keys (if the context argument is annotated as Context). I also moved some code to the module.

One nice thing about Context being (faked as) a TypedDict is, all existing def execute(self, context: Dict): annotations now suddenly become correct, because a TypedDict is automatically Dict[str, Any]! Everything works out on their own.

[potiuk]

Cool!

[github-actions]

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest main at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.

[ferruzzi]

Very cool. That was definitely a better solution than what I was doing, I'll drop mine and pick up something else to work on. For my own benefit, is there a particular reason why you dropped the context module in utils instead of in models?

[uranusjr]

Not really. I actually thought a while trying to decide, but eventually just picked one arbitrarily since it doesn't really matter either way.