82 lines
2.0 KiB
Python
82 lines
2.0 KiB
Python
|
from mypy.plugin import Plugin, FunctionContext
|
||
|
|
from mypy.types import Type, UnionType
|
||
|
|
|
||
|
|
|
||
|
|
class CadqueryPlugin(Plugin):
|
||
|
|
def get_function_hook(self, fullname: str):
|
||
|
|
|
||
|
|
if fullname == "cadquery.occ_impl.shapes._get":
|
||
|
|
|
||
|
|
return hook__get
|
||
|
|
|
||
|
|
elif fullname == "cadquery.occ_impl.shapes._get_one":
|
||
|
|
|
||
|
|
return hook__get_one
|
||
|
|
|
||
|
|
elif fullname == "cadquery.occ_impl.shapes._get_edges":
|
||
|
|
|
||
|
|
return hook__get_edges
|
||
|
|
|
||
|
|
elif fullname == "cadquery.occ_impl.shapes._get_wires":
|
||
|
|
|
||
|
|
return hook__get_wires
|
||
|
|
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def hook__get(ctx: FunctionContext) -> Type:
|
||
|
|
"""
|
||
|
|
Hook for cq.occ_impl.shapes._get
|
||
|
|
|
||
|
|
Based on the second argument values it adjusts return type to an Iterator of specific subclasses of Shape.
|
||
|
|
"""
|
||
|
|
|
||
|
|
if hasattr(ctx.args[1][0], "items"):
|
||
|
|
return_type_names = [el.value for el in ctx.args[1][0].items]
|
||
|
|
else:
|
||
|
|
return_type_names = [ctx.args[1][0].value]
|
||
|
|
|
||
|
|
return_types = UnionType([ctx.api.named_type(n) for n in return_type_names])
|
||
|
|
|
||
|
|
return ctx.api.named_generic_type("typing.Iterable", [return_types])
|
||
|
|
|
||
|
|
|
||
|
|
def hook__get_one(ctx: FunctionContext) -> Type:
|
||
|
|
"""
|
||
|
|
Hook for cq.occ_impl.shapes._get_one
|
||
|
|
|
||
|
|
Based on the second argument values it adjusts return type to a Union of specific subclasses of Shape.
|
||
|
|
"""
|
||
|
|
|
||
|
|
if hasattr(ctx.args[1][0], "items"):
|
||
|
|
return_type_names = [el.value for el in ctx.args[1][0].items]
|
||
|
|
else:
|
||
|
|
return_type_names = [ctx.args[1][0].value]
|
||
|
|
|
||
|
|
return UnionType([ctx.api.named_type(n) for n in return_type_names])
|
||
|
|
|
||
|
|
|
||
|
|
def hook__get_wires(ctx: FunctionContext) -> Type:
|
||
|
|
"""
|
||
|
|
Hook for cq.occ_impl.shapes._get_wires
|
||
|
|
"""
|
||
|
|
|
||
|
|
return_type = ctx.api.named_type("Wire")
|
||
|
|
|
||
|
|
return ctx.api.named_generic_type("typing.Iterable", [return_type])
|
||
|
|
|
||
|
|
|
||
|
|
def hook__get_edges(ctx: FunctionContext) -> Type:
|
||
|
|
"""
|
||
|
|
Hook for cq.occ_impl.shapes._get_edges
|
||
|
|
"""
|
||
|
|
|
||
|
|
return_type = ctx.api.named_type("Edge")
|
||
|
|
|
||
|
|
return ctx.api.named_generic_type("typing.Iterable", [return_type])
|
||
|
|
|
||
|
|
|
||
|
|
def plugin(version: str):
|
||
|
|
|
||
|
|
return CadqueryPlugin
|