diff --git a/mypy.ini b/mypy.ini new file mode 100644 index 000000000000..13deb9e234fb --- /dev/null +++ b/mypy.ini @@ -0,0 +1,2 @@ +[mypy] +plugins = compile_filename.plugin diff --git a/plugins/__init__.py b/plugins/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/plugins/compile_filename/__init__.py b/plugins/compile_filename/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/plugins/compile_filename/plugin.py b/plugins/compile_filename/plugin.py new file mode 100644 index 000000000000..ea780202355b --- /dev/null +++ b/plugins/compile_filename/plugin.py @@ -0,0 +1,32 @@ +from mypy.plugin import FunctionContext, Plugin +from mypy.types import Instance, Type + + +class CompileFilenamePlugin(Plugin): + def get_function_hook(self, fullname: str): + # Hook only the built-in compile function + if fullname == "builtins.compile": + return compile_hook + return None + + +def compile_hook(ctx: FunctionContext) -> Type: + # Arguments to compile: source, filename, mode, ... + # filename is arg index 1 (zero-based) + if len(ctx.arg_types) > 1 and ctx.arg_types[1]: + filename_type = ctx.arg_types[1][0] # first argument passed for filename param + if is_bytearray_type(filename_type): + ctx.api.fail("Passing 'bytearray' as filename to 'compile()' is not allowed", ctx.args[1][0]) + return ctx.default_return_type + + +def is_bytearray_type(typ: Type) -> bool: + # Check if the type is exactly bytearray + if isinstance(typ, Instance): + # The full name for builtins.bytearray is 'builtins.bytearray' + return typ.type.fullname == "builtins.bytearray" + return False + + +def plugin(version: str): + return CompileFilenamePlugin