Include the Function or Method context in the AST Node environment (#549)

This PR is a set of otherwise-independent changes I pulled out of #547. It includes the following changes:
 * The titular Function/Method context stored in the `NodeEnv` object of Lisp AST nodes
 * Changing the generated module-wide `__NS` variable to `_NS` because I learned that variables prefixed by double underscores are [always mangled](https://bugs.python.org/issue27793) in a class lexical context in Python regardless of the presence of a `global` statement.
 * Added a few tests for asserting the resolution of `def`'ed Basilisp Vars match the expected behavior from Clojure

Fixes #548

Author: chrisrink10

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
  * Added support for calling Python functions and methods with keyword arguments (#531)
  * Added support for Lisp functions being called with keyword arguments (#528)
  * Added support for multi-arity methods on `deftype`s (#534)
+ * Added metadata about the function or method context of a Lisp AST node in the `NodeEnv` (#548)
 
 ### Fixed
  * Fixed a bug where the Basilisp AST nodes for return values of `deftype` members could be marked as _statements_ rather than _expressions_, resulting in an incorrect `nil` return (#523)

src/basilisp/lang/compiler/analyzer.py

@@ -87,6 +87,7 @@
     Do,
     Fn,
     FnArity,
+    FunctionContext,
     HostCall,
     HostField,
     If,
@@ -309,7 +310,7 @@ def __init__(
     ) -> None:
         self._allow_unresolved_symbols = allow_unresolved_symbols
         self._filename = Maybe(filename).or_else_get(DEFAULT_COMPILER_FILE_PATH)
-        self._func_ctx: Deque[bool] = collections.deque([])
+        self._func_ctx: Deque[FunctionContext] = collections.deque([])
         self._is_quoted: Deque[bool] = collections.deque([])
         self._macro_ns: Deque[Optional[runtime.Namespace]] = collections.deque([])
         self._opts = (
@@ -406,34 +407,33 @@ def should_macroexpand(self) -> bool:
         return self._should_macroexpand
 
     @property
-    def is_async_ctx(self) -> bool:
-        """If True, the current node appears inside of an async function definition.
+    def func_ctx(self) -> Optional[FunctionContext]:
+        """Return the current function or method context of the current node, if one.
+        Return None otherwise.
+
         It is possible that the current function is defined inside other functions,
         so this does not imply anything about the nesting level of the current node."""
         try:
-            return self._func_ctx[-1] is True
+            return self._func_ctx[-1]
         except IndexError:
-            return False
+            return None
 
     @property
-    def in_func_ctx(self) -> bool:
-        """If True, the current node appears inside of a function definition.
+    def is_async_ctx(self) -> bool:
+        """Return True if the current node appears inside of an async function
+        definition. Return False otherwise.
+
         It is possible that the current function is defined inside other functions,
         so this does not imply anything about the nesting level of the current node."""
-        try:
-            self._func_ctx[-1]
-        except IndexError:
-            return False
-        else:
-            return True
+        return self.func_ctx == FunctionContext.ASYNC_FUNCTION
 
     @contextlib.contextmanager
-    def new_func_ctx(self, is_async: bool = False):
-        """Context manager which can be used to set a function context for child
-        nodes to examine. A new function context is pushed onto the stack each time
-        the Analyzer finds a new function definition, so there may be many nested
-        function contexts."""
-        self._func_ctx.append(is_async)
+    def new_func_ctx(self, context_type: FunctionContext):
+        """Context manager which can be used to set a function or method context for
+        child nodes to examine. A new function context is pushed onto the stack each
+        time the Analyzer finds a new function or method definition, so there may be
+        many nested function contexts."""
+        self._func_ctx.append(context_type)
         yield
         self._func_ctx.pop()
 
@@ -578,8 +578,14 @@ def syntax_position(self) -> NodeSyntacticPosition:
         parent node."""
         return self._syntax_pos[-1]
 
-    def get_node_env(self, pos: Optional[NodeSyntacticPosition] = None):
-        return NodeEnv(ns=self.current_ns, file=self.filename, pos=pos)
+    def get_node_env(self, pos: Optional[NodeSyntacticPosition] = None) -> NodeEnv:
+        """Return the current Node environment.
+
+        If a synax position is given, it will be included in the environment.
+        Otherwise, the position will be set to None."""
+        return NodeEnv(
+            ns=self.current_ns, file=self.filename, pos=pos, func_ctx=self.func_ctx
+        )
 
 
 MetaGetter = Callable[[Union[IMeta, Var]], bool]
@@ -810,10 +816,11 @@ def _def_ast(  # pylint: disable=too-many-branches,too-many-locals
             f"def names must be symbols, not {type(name)}", form=name
         )
 
+    children: vec.Vector[kw.Keyword]
     if nelems == 2:
         init = None
         doc = None
-        children: vec.Vector[kw.Keyword] = vec.Vector.empty()
+        children = vec.Vector.empty()
     elif nelems == 3:
         with ctx.expr_pos():
             init = _analyze_form(ctx, runtime.nth(form, 2))
@@ -895,7 +902,6 @@ def _def_ast(  # pylint: disable=too-many-branches,too-many-locals
         var=var,
         init=init,
         doc=doc,
-        in_func_ctx=ctx.in_func_ctx,
         children=children,
         env=def_node_env,
     )
@@ -1037,7 +1043,7 @@ def __deftype_classmethod(
         has_vargs, fixed_arity, param_nodes = __deftype_method_param_bindings(
             ctx, params
         )
-        with ctx.expr_pos():
+        with ctx.new_func_ctx(FunctionContext.CLASSMETHOD), ctx.expr_pos():
             stmts, ret = _body_ast(ctx, runtime.nthrest(form, 2))
         method = DefTypeClassMethodArity(
             form=form,
@@ -1097,7 +1103,7 @@ def __deftype_method(
 
         loop_id = genname(method_name)
         with ctx.new_recur_point(loop_id, param_nodes):
-            with ctx.expr_pos():
+            with ctx.new_func_ctx(FunctionContext.METHOD), ctx.expr_pos():
                 stmts, ret = _body_ast(ctx, runtime.nthrest(form, 2))
             method = DefTypeMethodArity(
                 form=form,
@@ -1160,7 +1166,7 @@ def __deftype_property(
 
         assert not has_vargs, "deftype* properties may not have arguments"
 
-        with ctx.expr_pos():
+        with ctx.new_func_ctx(FunctionContext.PROPERTY), ctx.expr_pos():
             stmts, ret = _body_ast(ctx, runtime.nthrest(form, 2))
         prop = DefTypeProperty(
             form=form,
@@ -1192,7 +1198,7 @@ def __deftype_staticmethod(
     """Emit a node for a :staticmethod member of a deftype* form."""
     with ctx.hide_parent_symbol_table(), ctx.new_symbol_table(method_name):
         has_vargs, fixed_arity, param_nodes = __deftype_method_param_bindings(ctx, args)
-        with ctx.expr_pos():
+        with ctx.new_func_ctx(FunctionContext.STATICMETHOD), ctx.expr_pos():
             stmts, ret = _body_ast(ctx, runtime.nthrest(form, 2))
         method = DefTypeStaticMethodArity(
             form=form,
@@ -1652,7 +1658,10 @@ def _do_ast(ctx: AnalyzerContext, form: ISeq) -> Do:
 
 
 def __fn_method_ast(  # pylint: disable=too-many-branches,too-many-locals
-    ctx: AnalyzerContext, form: ISeq, fnname: Optional[sym.Symbol] = None
+    ctx: AnalyzerContext,
+    form: ISeq,
+    fnname: Optional[sym.Symbol] = None,
+    is_async: bool = False,
 ) -> FnArity:
     with ctx.new_symbol_table("fn-method"):
         params = form.first
@@ -1711,7 +1720,9 @@ def __fn_method_ast(  # pylint: disable=too-many-branches,too-many-locals
 
         fn_loop_id = genname("fn_arity" if fnname is None else fnname.name)
         with ctx.new_recur_point(fn_loop_id, param_nodes):
-            with ctx.expr_pos():
+            with ctx.new_func_ctx(
+                FunctionContext.ASYNC_FUNCTION if is_async else FunctionContext.FUNCTION
+            ), ctx.expr_pos():
                 stmts, ret = _body_ast(ctx, form.rest)
             method = FnArity(
                 form=form,
@@ -1770,11 +1781,11 @@ def _fn_ast(  # pylint: disable=too-many-branches
                 form=form,
             )
 
+        name_node: Optional[Binding]
         if isinstance(name, sym.Symbol):
-            name_node: Optional[Binding] = Binding(
+            name_node = Binding(
                 form=name, name=name.name, local=LocalType.FN, env=ctx.get_node_env()
             )
-            assert name_node is not None
             is_async = _is_async(name) or isinstance(form, IMeta) and _is_async(form)
             kwarg_support = (
                 __fn_kwargs_support(name)
@@ -1802,23 +1813,24 @@ def _fn_ast(  # pylint: disable=too-many-branches
                 form=form,
             )
 
-        with ctx.new_func_ctx(is_async=is_async):
-            if isinstance(arity_or_args, llist.List):
-                arities = vec.vector(
-                    map(
-                        partial(__fn_method_ast, ctx, fnname=name),
-                        runtime.nthrest(form, idx),
-                    )
-                )
-            elif isinstance(arity_or_args, vec.Vector):
-                arities = vec.v(
-                    __fn_method_ast(ctx, runtime.nthrest(form, idx), fnname=name)
+        if isinstance(arity_or_args, llist.List):
+            arities = vec.vector(
+                map(
+                    partial(__fn_method_ast, ctx, fnname=name, is_async=is_async),
+                    runtime.nthrest(form, idx),
                 )
-            else:
-                raise AnalyzerException(
-                    "fn form must match: (fn* name? [arg*] body*) or (fn* name? method*)",
-                    form=form,
+            )
+        elif isinstance(arity_or_args, vec.Vector):
+            arities = vec.v(
+                __fn_method_ast(
+                    ctx, runtime.nthrest(form, idx), fnname=name, is_async=is_async
                 )
+            )
+        else:
+            raise AnalyzerException(
+                "fn form must match: (fn* name? [arg*] body*) or (fn* name? method*)",
+                form=form,
+            )
 
         nmethods = count(arities)
         assert nmethods > 0, "fn must have at least one arity"

src/basilisp/lang/compiler/generator.py

@@ -131,7 +131,7 @@
 _SET_BANG_TEMP_PREFIX = "set_bang_val"
 _THROW_PREFIX = "lisp_throw"
 _TRY_PREFIX = "lisp_try"
-_NS_VAR = "__NS"
+_NS_VAR = "_NS"
 
 
 GeneratorException = partial(CompilerException, phase=CompilerPhase.CODE_GENERATION)
@@ -750,15 +750,19 @@ def _def_to_py_ast(  # pylint: disable=too-many-branches
             # complaining that we assign the value prior to global declaration.
             def_dependencies = list(
                 chain(
-                    [ast.Global(names=[safe_name])] if node.in_func_ctx else [],
+                    [ast.Global(names=[safe_name])]
+                    if node.env.func_ctx is not None
+                    else [],
                     def_ast.dependencies,
                 )
             )
         else:
             def_dependencies = list(
                 chain(
                     def_ast.dependencies,
-                    [ast.Global(names=[safe_name])] if node.in_func_ctx else [],
+                    [ast.Global(names=[safe_name])]
+                    if node.env.func_ctx is not None
+                    else [],
                     [
                         ast.Assign(
                             targets=[ast.Name(id=safe_name, ctx=ast.Store())],
@@ -775,7 +779,9 @@ def _def_to_py_ast(  # pylint: disable=too-many-branches
         # root.
         func = _INTERN_UNBOUND_VAR_FN_NAME
         extra_args = []
-        def_dependencies = [ast.Global(names=[safe_name])] if node.in_func_ctx else []
+        def_dependencies = (
+            [ast.Global(names=[safe_name])] if node.env.func_ctx is not None else []
+        )
 
     meta_ast = gen_py_ast(ctx, node.meta)
 

src/basilisp/lang/compiler/nodes.py

@@ -272,6 +272,15 @@ class ConstType(Enum):
 LoopID = str
 
 
+class FunctionContext(Enum):
+    FUNCTION = kw.keyword("function")
+    ASYNC_FUNCTION = kw.keyword("async-function")
+    METHOD = kw.keyword("method")
+    CLASSMETHOD = kw.keyword("classmethod")
+    STATICMETHOD = kw.keyword("staticmethod")
+    PROPERTY = kw.keyword("property")
+
+
 class KeywordArgSupport(Enum):
     APPLY_KWARGS = kw.keyword("apply")
     COLLECT_KWARGS = kw.keyword("collect")
@@ -296,6 +305,7 @@ class NodeEnv:
     line: Optional[int] = None
     col: Optional[int] = None
     pos: Optional[NodeSyntacticPosition] = None
+    func_ctx: Optional[FunctionContext] = None
 
 
 @attr.s(auto_attribs=True, frozen=True, slots=True)
@@ -360,7 +370,6 @@ class Def(Node[SpecialForm]):
     var: Var
     init: Optional[Node]
     doc: Optional[str]
-    in_func_ctx: bool
     env: NodeEnv
     meta: NodeMeta = None
     children: Sequence[kw.Keyword] = vec.Vector.empty()

src/basilisp/lang/runtime.py

@@ -1656,13 +1656,13 @@ def init_ns_var() -> Var:
 def bootstrap_core(compiler_opts: CompilerOpts) -> None:
     """Bootstrap the environment with functions that are either difficult to express
     with the very minimal Lisp environment or which are expected by the compiler."""
-    __NS = Maybe(Var.find(NS_VAR_SYM)).or_else_raise(
+    _NS = Maybe(Var.find(NS_VAR_SYM)).or_else_raise(
         lambda: RuntimeException(f"Dynamic Var {NS_VAR_SYM} not bound!")
     )
 
     def in_ns(s: sym.Symbol):
         ns = Namespace.get_or_create(s)
-        __NS.value = ns
+        _NS.value = ns
         return ns
 
     # Vars used in bootstrapping the runtime

tests/basilisp/compiler_test.py

@@ -4264,6 +4264,66 @@ def test_nested_bare_sym_will_not_resolve(self, lcompile: CompileFn):
         with pytest.raises(compiler.CompilerException):
             lcompile("basilisp.lang.map.MapEntry.of")
 
+    @pytest.mark.parametrize(
+        "code",
+        [
+            """
+            (fn []
+              (def a :a)
+              (var a))""",
+            """
+            (import* collections.abc)
+            (deftype* Definer []
+              :implements [collections.abc/Callable]
+              (--call-- [this] (def a :a) (var a)))
+            (Definer)""",
+        ],
+    )
+    def test_symbol_deffed_in_fn_or_method_will_resolve_in_fn_or_method(
+        self, ns: runtime.Namespace, lcompile: CompileFn, code: str,
+    ):
+        # This behavior is peculiar and perhaps even _wrong_, but it matches how
+        # Clojure treats Vars defined in functions. Of course, generally speaking,
+        # Vars should not be defined like this so I suppose it's not a huge deal.
+        fn = lcompile(code)
+
+        resolved_var = ns.find(sym.symbol("a"))
+        assert not resolved_var.is_bound
+
+        returned_var = fn()
+        assert returned_var is resolved_var
+        assert returned_var.is_bound
+        assert returned_var.value == kw.keyword("a")
+
+    @pytest.mark.parametrize(
+        "code",
+        [
+            """
+            (do
+              (fn [] (def a :a))
+              (var a))""",
+            """
+            (fn [] (def a :a))
+            (var a)""",
+            """
+            (import* collections.abc)
+            (deftype* Definer []
+              :implements [collections.abc/Callable]
+              (--call-- [this] (def a :a)))
+            (var a)""",
+        ],
+    )
+    def test_symbol_deffed_in_fn_or_method_will_resolve_outside_fn_or_method(
+        self, ns: runtime.Namespace, lcompile: CompileFn, code: str
+    ):
+        var = lcompile(code)
+        assert not var.is_bound
+
+        resolved_var = ns.find(sym.symbol("a"))
+        assert not resolved_var.is_bound
+
+        assert var is resolved_var
+
     def test_local_deftype_classmethod_resolves(self, lcompile: CompileFn):
         Point = lcompile(
             """