PyCharm code completion and dynamic casting

Question

I'm pretty new to Python 3 and type-hints and I'm wondering if it's possible to do this in PyCharm or other IDE's.

Example

Here's a simplified example of the code.

import dataclasses
from typing import Type

@dataclasses.dataclass()
class BasePlugin:
    version: str
    arguments: str = None

@dataclasses.dataclass()
class PythonPlugin(BasePlugin):
    version: str = 3.7
    script: str = None

@dataclasses.dataclass()
class MayaPlugin(BasePlugin):
    version: str = 2022
    scene_file: str = None
    script_job: str = None


class Job:
    def __init__(self, plugin: Type[BasePlugin]):
        self.options = plugin()


job_a = Job(plugin=PythonPlugin)

Ideally, I want to be able to create an instance of Job and use dot-notation to modify the options fields.
e.g. job_a.options.script = 'c:\test\do_stuff.py'

Problem

Incorrect result

With the input type-hint, code completion only shows the available attributes of the base class.

Code sample:

class Job:
    def __init__(self, plugin: Type[BasePlugin]):
        self.options = plugin()


job_a = Job(plugin=PythonPlugin)
job_a.options.

Result: Code completion with input type-hint

---

Incorrect result

Without the type-hint, code completion shows nothing, as expected.

Code sample:

class Job:
    def __init__(self, plugin):
        self.options = plugin()


job_a = Job(plugin=PythonPlugin)
job_a.options.

Result: Code completion without input type-hint

---

Desired result

What I want is the code completion for the specific subclass I'm providing on instantiation. In the example below, I'm overriding self.options with the plugin to show the desired result.

Code sample:

class Job:
    def __init__(self, plugin: Type[BasePlugin]):
        self.options = PythonPlugin()  # Overriding to simulate the desired result.


job_a = Job(plugin=PythonPlugin)
job_a.options.

Result: Code completion as desired

Solutions

1. Is it possible to achieve the desired result by passing in an argument?
2. If not, is my only other option to subclass the job types?

class Job:
    options: BasePlugin
    ...

class PythonJob(Job):
    options: PythonPlugin = PythonPlugin()
    ...

class MayaJob(Job):
    options: MayaPlugin = MayaPlugin()
    ...

3. Any other options/examples I need to consider?

Any help is much appreciated.

Answer 1

You can do this with TypeVars and Generics:

import dataclasses
from typing import Generic, Type, TypeVar

BP = TypeVar("BP")

@dataclasses.dataclass()
class BasePlugin:
    version: str
    arguments: str = None

@dataclasses.dataclass()
class PythonPlugin(BasePlugin):
    version: str = 3.7
    script: str = None

@dataclasses.dataclass()
class MayaPlugin(BasePlugin):
    version: str = 2022
    scene_file: str = None
    script_job: str = None


class Job(Generic[BP]):
    def __init__(self, plugin: Type[BP]):
        self.options: BP = plugin()


job_a = Job(plugin=PythonPlugin)
job_a.options.script  # auto-completes

job_b = Job(plugin=MayaPlugin)
job_b.options.scene_file  # auto-completes