rework and test custom builtins

pyflakes/checker.py

@@ -21,7 +21,21 @@
 
 PYPY = hasattr(sys, 'pypy_version_info')
 
-builtin_vars = dir(builtins)
+builtin_vars = frozenset(dir(builtins)) | {
+    # Globally defined names which are not attributes of the builtins module, or
+    # are only present on some platforms.
+    '__file__', '__builtins__', '__annotations__', 'WindowsError'
+}
+
+
+@functools.cache
+def _custom_builtins() -> frozenset[str]:
+    var = os.environ.get('PYFLAKES_BUILTINS')
+    if var is not None:
+        return frozenset(var.split(','))
+    else:
+        return frozenset()
+
 
 parse_format_string = string.Formatter().parse
 
@@ -579,11 +593,6 @@ class DetectClassScopedMagic:
     names = dir()
 
 
-# Globally defined names which are not attributes of the builtins module, or
-# are only present on some platforms.
-_MAGIC_GLOBALS = ['__file__', '__builtins__', '__annotations__', 'WindowsError']
-
-
 def getNodeName(node):
     # Returns node.id, or node.name, or None
     if hasattr(node, 'id'):     # One of the many nodes with an id
@@ -714,19 +723,14 @@ class Checker:
     offset = None
     _in_annotation = AnnotationState.NONE
 
-    builtIns = set(builtin_vars).union(_MAGIC_GLOBALS)
-    _customBuiltIns = os.environ.get('PYFLAKES_BUILTINS')
-    if _customBuiltIns:
-        builtIns.update(_customBuiltIns.split(','))
-    del _customBuiltIns
-
     def __init__(self, tree, filename='(none)', builtins=None,
                  withDoctest='PYFLAKES_DOCTEST' in os.environ):
         self._nodeHandlers = {}
         self._deferred = collections.deque()
         self.deadScopes = []
         self.messages = []
         self.filename = filename
+        self.builtIns = builtin_vars | _custom_builtins()
         if builtins:
             self.builtIns = self.builtIns.union(builtins)
         self.withDoctest = withDoctest

pyflakes/test/test_custom_builtins.py

@@ -0,0 +1,29 @@
+import contextlib
+import os
+from unittest import mock
+from pyflakes import messages as m
+from pyflakes.test.harness import TestCase
+from pyflakes.checker import _custom_builtins
+
+
+@contextlib.contextmanager
+def _clear_custom_builtins_cache():
+    _custom_builtins.cache_clear()
+    try:
+        yield
+    finally:
+        _custom_builtins.cache_clear()
+
+
+class TestCustomBuiltins(TestCase):
+    def test_custom_builtins_from_init(self):
+        self.flakes('unknown', m.UndefinedName)
+        self.flakes('unknown', builtins=('unknown',))
+
+    def test_custom_builtins_from_env(self):
+        self.flakes('y = a + b', m.UndefinedName, m.UndefinedName)
+        with (
+                _clear_custom_builtins_cache(),
+                mock.patch.dict(os.environ, {'PYFLAKES_BUILTINS': 'a,b'}),
+        ):
+            self.flakes('y = a + b')