Add type annotations to opengl_mobject.py

[RBerga06]

Overview: What does this pull request change?

Part of #3375.

Motivation and Explanation: Why and how do your changes improve the library?

This PR adds type annotations to manim/mobject/opengl/opengl_mobject.py, to improve type coverage for manim's public interface.

Edit: By thoroughly annotating a number of methods, I also discovered and fixed what appeared to be typos or mistakes. For reference, I'm listing them here, because these changes can in principle affect runtime behaviour.

• Fixed a typo in OpenGLMobject.restore (commit 40dc43f)
• Fixed extra **kwargs being passed to OpenGLMobject.apply_points_function that could lead to runtime TypeErrors (see my comment below)
• Implemented a missing OpenGLMobject.get_array_attrs() method (commit 2909380)
• Fixed a missing return self in OpenGLMobject.pointwise_become_partial (commit 5d5de9a)

Reviewer Checklist

• The PR title is descriptive enough for the changelog, and the PR is labeled correctly
• If applicable: newly added non-private functions and classes have a docstring including a short summary and a PARAMETERS section
• If applicable: newly added functions and classes are tested

[RBerga06]

I also made the manim.utils.config_ops._Data descriptor generic, to allow annotating the corresponding attributes in OpenGLMobject with the relevant array types. I also added a _HasData protocol to avoid using _Data in objects that are not guaranteed to have the data attribute.

[RBerga06]

I have also discovered and fixed a typo in OpenGLMobject.restore (commit 40dc43f).

[RBerga06]

It is not completely clear where extra **kwargs should be allowed (sometimes they are, sometimes they're not).

For example, I noticed that the apply_points_function method only accepts a fixed set of kwargs and that the scale method explicitely sets all of them. However, scale also forwards to apply_points_function additional **kwargs! This means that passing any extra keyword arguments to scale WILL introduce a runtime error.

That's why, at first, I annotated **kwargs in scale as **kwargs: Never: they must not exist at runtime. However, as soon as I did this mypy rejected my code, because it had introduced an error in manim/animation/growing.py, at line 208, where the scale method is invoked on a Mobject | OpenGLMobject with additional parameters (in particular, scale_tips=True), thus triggering a runtime error every time the object was an OpenGLObject!

I suspect the intention was to override the scale method in Mobject/OpenGLMobject subclasses and simply ignore all unrecognized keyword argument in the base class. Therefore, I kept **kwargs: object in scale and removed them from the apply_points_function call.

I think it would help to define a MobjectLike protocol with the interface that should be common to both Mobject and OpenGLMobject, but I think I should introduce it in a separate PR. This way, at least, the common object interface is clear and type checked.

@henrikmidtiby whay do you think about this proposal? I'm asking you because I saw you're working on typing Mobject in #4388 and I guess you might have noticed similar problems to the ones I'm describing here.

[RBerga06]

I also think clarifying the current state and intended Mobject/OpenGLObject interface/behaviour might also help the work in #3112.

[henrikmidtiby]

Wow a lot of good stuff is happening here!

Here comes a list of observations, that I have made about things that I would have done differently; which is not necessarily better...

• In some locations you have typed **kwargs: dict[str, object], I have used the following earlier **kwargs: Any.

• The use of Unpack[_Kw_Arrange] is new to me. Do you have any good description of this approach? It seems quite flexible to use.

• There is a number of # type: ignore[...] statements in the code. I prefer to insert the triggered error in the line above where the error was triggered. The intention is to make it easier to get input from other on how to fix the issue.

• In some cases you get the error error: Name "..." already defined on line ... [no-redef]. My approach to deal with these issues is typically to introduce a new variable. At the moment mypy is not able to deal with variables that changes type inside a function or method.

• An approach to the error error: Call to untyped function "ShaderWrapper" in typed context [no-untyped-call] is to locate the untyped function and just add type annotations to that in the other file, without activating type checking of that file in mypy.ini.

[RBerga06]

Regarding the first two points, the Unpack[TD] (where TD is a TypedDict) idiom has been introduced fairly recently (the typing spec seems to describe it pretty well). In essence, because **kwargs: T binds the keyword arguments to a dict[str, T], now that we have a way of typing {"a": 3, "b": "foo"} more strictly than dict[str, int | str] (i.e. we have TypedDicts), we can also better annotate those **kwargs. You can think of **kwargs: Unpack[TypedDict] as the moral equivalent of *args: *TupleType (where TupleType is a tuple[<something...>]) - in fact, *TupleType is technically a shorthand for Unpack[TupleType]. Defining the TypedDicts with total=False makes all defined **kwargs optional.

If at some point I've used **kwargs: dict[str, object], I think that's an unintended mistake; what I really meant was **kwargs: object, as a way to communicate that the function accepts additional keyword arguments, but does not care about them / handle them in any way. But now that I think about it, maybe that's wrong. Consider this case:

class A:
    def foo(self, **kwargs: object) -> None: ...

class B(A):
    def foo(self, x: int = 3, **kwargs: object) -> None: ...

def use_a(a: A):
    a.foo(x="hello")  # ok because "hello": str <: object

use_a(A())  # all fine
use_a(B())  # TypeError at runtime!

In fact, type checkers should (and maybe do?) reject B's definition - I admit I have not really checked it yet Edit: mypy rejects it but strangely enough pyright allows it (even in strict mode!).

So yes, I think you're right: I should've used **kwargs: Any in that situation.

This reminds me I still have to check if extra **kwargs are allowed in functions annotated with **kwargs: Unpack[<some TypedDict>] - I suspect they're not, but I think there is a way of allowing extra members on TypedDict. I'll report back very soon.

Edit: they're not supported, but they might will be soon: see the draft PEP for that. pyright already has implemented the PEP (although support is still marked as experimental) and mypy will follow.

I've opened #3375 to track this problem (among other things). but because the PEP is still a draft, I think we should just use **kwargs: Any everywhere for now.

[RBerga06]

By the way, I have started using TypeDicts to annotate **kwargs in the first place because I remember myself, at some point in the past, trying to desperately figure out which color-(or maybe shade-)related arguments were supported by a specific Mobject subclass (and which types were allowed). I had to follow all **kwargs in __init__ across two or three classes before finding the related documentation. I have also introduced a couple of docstrings in _Kw_OpenGLMobject, in the hope that they are picked up by my IDE (basedpyright language server in VSCodium).

[RBerga06]

• There is a number of # type: ignore[...] statements in the code. I prefer to insert the triggered error in the line above where the error was triggered. The intention is to make it easier to get input from other on how to fix the issue.

Yes, that actually makes sense. In a couple of places, I think I have tried my best to explain why the # type: ignore was needed.

I'll try and explain the rest of them; although, if I recall correctly, they're all about incompatible assignment between Mobject and OpenGLMobject. I don't fully understand why at some point we're creating a Line (definitely a Mobject, right?) and treating it like a OpenGLMobject, but a part from that it seems like non-manim.mobject.* APIs accept Mobject when they really should be accepting Mobject | OpenGLMobject.

Ah yes, there was also a very long function where we were shadowing function arguments with new variables in the function body and mypy seemed to keep the old type (giving rise to “str is not callable” and similar error messages). When I have time, I'll try and submit a mypy issue for that.

As a side note, I still cannot figure out how to get mypy correctly configured (I've always been a pyright/basedpyright user). I installed mypy with uv, directly in the dev environment (without committing changes to pyproject.toml/uv.lock) but mypy still reveals more errors than the pre-commit ci. Maybe I'll investigate a little bit more; for the time being, I've disabled the IDE extension and I'm only relying on pre-commit CI (and basedpyright).

[RBerga06]

Ah, missed this part of your comment:

In some cases you get the error error: Name "..." already defined on line ... [no-redef]. My approach to deal with these issues is typically to introduce a new variable. At the moment mypy is not able to deal with variables that changes type inside a function or method.

Yes, I agree it would be a solution; however, I think in some cases (like this very very long function) it pollutes the namespace and maybe harms code readibility. It also requires updating absolutely all references without missing anything, and has therefore to be reviewed with extreme care.

In languages like Rust or Swift you can clearly shadow variables with a redeclaration (let <name> or let mut <name>/var <name>) and, in typed Python, the most similar construct is a variable annotation (<name>: <type>). That's why I think long-term that should be the correct solution.

Because type checkers like mypy are evidently trying to support this kind of pattern (see --allow-redefinition-new in the mypy docs), I think short-term it's ok to keep the # type: ignore[no-redef], until the feature is fully implemented in mypy (and maybe even defined in the typing spec).

Also, I don't understand why mypy is still complaining about it, even if I have enabled the stable allow_redefinition = True in mypy.ini - maybe it's a mypy bug? 🤔

Anyways, if you still think we should use different names, I don't feel too strongly about this: I'll update the PR accordingly.

[RBerga06]

An approach to the error error: Call to untyped function "ShaderWrapper" in typed context [no-untyped-call] is to locate the untyped function and just add type annotations to that in the other file, without activating type checking of that file in mypy.ini.

Yes, I thought about that: I was trying to keep the changes as self-contained as possible. But it doesn't look like it should be too hard, I'll try and annotate those methods asap. Edit: done

By the way, thank you @henrikmidtiby for the review!

[RBerga06]

Ah, one last thing: if you reckon that the **kwargs in OpenGLMobject and its methods have a non-empty intersection with the ones in Mobject, I'd be happy to submit a PR where we store the definitions of the common ones in a shared location, maybe some new manim.mobject._abc or manim.mobject._interface module/subpackage, per my idea in this previous comment.

[henrikmidtiby]

@RBerga06 I am happy to provide feedback.

I really like the use use of Unpack[TypedDict] for adding type annotations to **kwargs. This approach is much better than the **kwargs: Any I have used until now.

I'll try and explain the rest of them; although, if I recall correctly, they're
all about incompatible assignment between Mobject and OpenGLMobject.
Yes, that can be quite difficult to get right.

Regarding the issue Name ... already defined on line ... [no-redef] with the method arrange_in_grid of 270 lines in opengl_mobject.py. To me the main issue is that the method needs to be refactored, right now way to much is happening inside the method.

I get your point that the type checker should be able to deal with the situation, this is however not the case yet ... For me the cost of polluting the namespace with one additional variable, is worth it, if we can get rid of any type: ignore statements.

Regarding to locate a common subset of **kwargs between Mobject and OpenGLMobject, I would like to raise the question on the discord server.

Finally the size of this PR is getting quite large with multiple changes at once (e.g. add type annotations and update dependency on typing-extensions). It is much easier for the reviewers to handle several small PR than one large PR.

@chopan050 Would you have time to look at this PR?

[RBerga06]

@henrikmidtiby

Finally the size of this PR is getting quite large with multiple changes at once (e.g. add type annotations and update dependency on typing-extensions).

Yes, I agree this is a pretty large PR. Given that **kwargs: Unpack[...] annotations require both adding several TypedDicts and bumping typing-extensions and given that we might even be moving those TypedDicts around anyway (to also type OpenGLMobject), I think the right thing to do is probably remove those changes from this PR (this ways it only adds type annotations) and reintroduce those changes in a dedicated PR, once we've discussed the design on Discord.

I'll update this PR accordingly ASAP (thus annotating **kwargs: Any for now). Edit: done.