subclassing Pyomo model for auto-completion / type hinting

Question

I'm considering implementing a schema as below for one of my larger pyomo projects. (Yes, I'm intentionally using m instead of the ubiquitous self for consistency in model.) In the project the model creation / solving / post-processing are all separate modules and I'd like to be able to employ the IDE's auto-complete and type hinting on the model, which would make it quicker / more accurate to locate key variables, sets, etc. in the post-processing and tidy up the project.

I'm considering sub-classing ConcreteModel to do this.

Is this a good/horrible idea?

Are there other options for achieving this?

I've seen this posting which indicates sub-classing has been considered for other reasons.

import pyomo.environ as pyo


class MyModel(pyo.ConcreteModel):
    def __init__(m):
        super().__init__()
        m.S = pyo.Set(initialize=[1, 2, 3])
        m.x = pyo.Var(m.S)

        @m.Constraint(m.S)
        def upper_limit(m, s):
            return m.x[s] <= 5

        # ... etc.


m1 = MyModel()

# meanwhile.... somewhere else in post-processing:
def list_set_vals(m: MyModel) -> list[int]:
    return list(m.S)  # <-- auto completion available

print(list_set_vals(m1))

Answer 1

I don't see any obvious pitfalls with doing that (and autocompletion would certainly be a nice benefit). ConcreteModel is really just semantic sugar on top of Block, and as long as you don't change the class / instance ctype (it should still return Block) then things like the solvers and transformations should still all work the same.

Other projects have made extensive use of subclassing Block (e.g., IDAES), although there the goal is to also support indexed subclasses, so the infrastructure to effect both scalar and indexed components is a little more complex.

Separately, Pyomo is way overdue on reorganizing / expanding the standard examples. This use case strikes me as a nice one to add to the standard set -- please consider submitting it to the project.