Dialect Group

def __call__(
    self,
    py_func: Callable[Param, RetType] | None = None,
    *args: PassParams.args,
    **options: PassParams.kwargs,
) -> Method[Param, RetType] | MethodTransform[Param, RetType]:
    """create a method from the python function.

    Args:
        py_func (Callable): the python function to create the method from.
        args (PassParams.args): the arguments to pass to the run_pass function.
        options (PassParams.kwargs): the keyword arguments to pass to the run_pass function.

    Returns:
        Method: the method created from the python function.
    """
    from kirin.lowering import Lowering

    emit_ir = Lowering(self)

    def wrapper(py_func: Callable) -> Method:
        if py_func.__name__ == "<lambda>":
            raise ValueError("Cannot compile lambda functions")

        lineno_offset, file = 0, ""
        frame = inspect.currentframe()
        if frame and frame.f_back is not None and frame.f_back.f_back is not None:
            call_site_frame = frame.f_back.f_back
            if py_func.__name__ in call_site_frame.f_locals:
                raise CompilerError(
                    f"overwriting function definition of `{py_func.__name__}`"
                )

            lineno_offset = call_site_frame.f_lineno - 1
            file = call_site_frame.f_code.co_filename

        code = emit_ir.run(py_func, lineno_offset=lineno_offset)
        mt = Method(
            mod=inspect.getmodule(py_func),
            py_func=py_func,
            sym_name=py_func.__name__,
            arg_names=["#self#"] + inspect.getfullargspec(py_func).args,
            dialects=self,
            code=code,
            file=file,
        )
        if doc := inspect.getdoc(py_func):
            mt.__doc__ = doc

        if self.run_pass is not None:
            self.run_pass(mt, *args, **options)
        return mt

    if py_func is not None:
        return wrapper(py_func)
    return wrapper

def __contains__(self, dialect) -> bool:
    """check if the dialect is in the group.

    Args:
        dialect (Union[Dialect, ModuleType]): the dialect to check.

    Returns:
        bool: True if the dialect is in the group, False otherwise.
    """
    return self.map_module(dialect) in self.data

def add(self, dialect: Union["Dialect", ModuleType]) -> "DialectGroup":
    """add a dialect to the group.

    Args:
        dialect (Union[Dialect, ModuleType]): the dialect to add

    Returns:
        DialectGroup: the new dialect group with the added
    """
    return self.union([dialect])

def discard(self, dialect: Union["Dialect", ModuleType]) -> "DialectGroup":
    """discard a dialect from the group.

    !!! note
        This does not raise an error if the dialect is not in the group.

    Args:
        dialect (Union[Dialect, ModuleType]): the dialect to discard

    Returns:
        DialectGroup: the new dialect group with the discarded dialect.
    """
    dialect_ = self.map_module(dialect)
    return DialectGroup(
        dialects=frozenset(
            each for each in self.data if each.name != dialect_.name
        ),
        run_pass=self.run_pass_gen,  # pass the run_pass_gen function
    )

@staticmethod
def map_module(dialect: Union["Dialect", ModuleType]) -> "Dialect":
    """map the module to the dialect if it is a module.
    It assumes that the module has a `dialect` attribute
    that is an instance of [`Dialect`][kirin.ir.Dialect].
    """
    if isinstance(dialect, ModuleType):
        return getattr(dialect, "dialect")
    return dialect

def union(self, dialect: Iterable[Union["Dialect", ModuleType]]) -> "DialectGroup":
    """union a set of dialects to the group.

    Args:
        dialect (Iterable[Union[Dialect, ModuleType]]): the dialects to union

    Returns:
        DialectGroup: the new dialect group with the union.
    """
    return DialectGroup(
        dialects=self.data.union(frozenset(self.map_module(d) for d in dialect)),
        run_pass=self.run_pass_gen,  # pass the run_pass_gen function
    )