python ensureimport but within TYPE_CHECKING block

[Distortedlogic]

I am working on using grit to locally monkey patch polars in my .venv to get IDE support for my plugins since polars does not provide a way to do this yet and the only solution I found so far is just not de wey. I use ensureimport but I need to ensure the import is in a if TYPE_CHECKING: to avoid circular imports.

References:
pola-rs/polars#14475
https://github.com/StefanBRas/polugins

Big fan of grit by the way, I had been looking for exactly this for a long time.

I think above is pretty straight forward on the core issue but I am providing more info below just in case for any further comment on my approach (is a work in progress, and is my first use of grit) .. Any additional thoughts etc are welcome.

Resources:

engine marzano(0.1)
language python

pattern extract_polats_plugins(){
    decorated_definition($decorators, $definition) where {
        $decorators <: some `pl.api.register_lazyframe_namespace("aggrid")`,
        // $definition <: class_definition($name, $superclasses, $body) where {
        //     $body <: contains function_definition($name,$parameters) where {
        //         $name <: "__init__",
        //         $parameters <: some typed_parameter(type=`pl.$pl_class_name`)
        //     }
        // }
        $definition => `$attr_name`
    }
}

pattern monkey_patch_polars_class($module_path, $class_name, $attr_name){
    class_definition($name, $superclasses, body=$class_body) as $cls where  {
        $name <: or { "LazyFrame", "DataFrame", "Expr", "Series" },
        $class_name <: ensure_import_from(source=$module_path),
        !$class_body <: contains bubble `$attr_name:$module_path`,
        $class_body => `$class_body\n$attr_name:$module_path`
  }
}

sequential{
    maybe bubble($body) file($body) where $body <: contains extract_polats_plugins(),
    maybe bubble($body) file($body) where $body <: contains monkey_patch_polars_class(module_path="src.packages.lm_polars.aggrid",class_name="Aggrid",attr_name="agrid"),
}

[morgante]

I'm glad you're finding Grit useful!

ensure_imported_from is actually just part of the standard library, so you could make a copy of that file that wraps the import inside an if TYPE_CHECKING: block when inserting the import. Feel free to hop into Discord if you have more questions.

For your pattern itself, you likely don't need $body to pierce the bubble. I'm also not sure if you really need to use sequential - the purpose of sequential is to run one pattern after the other (ie. the output of the first pattern is necessary for the second).. If you just need to run two patterns on each file you can do it like this:

file($body) where {
  $body <: contains any {
    extract_polats_plugins(),
    monkey_patch_polars_class(module_path="src.packages.lm_polars.aggrid",class_name="Aggrid",attr_name="agrid"),
  }
}

[Distortedlogic]

ill join discord cause i do have questions and have a lot of AST stuff I already do that grit would make a ton ton better, plus more ast stuff I want to do.

the use of sequential is to prepare to use multifile (i didnt see/find anything about the studio supporting multifile to dev on). My plugins will be in several different files, and the polars package file will be inside my venv where its installed.