From 4ef3678a28f3521572653b38974999bf6b2778cf Mon Sep 17 00:00:00 2001
From: Hendrik van Antwerpen <hendrikvanantwerpen@github.com>
Date: Mon, 22 Apr 2024 11:33:16 +0200
Subject: [PATCH] Add support for lambda parameters

---
 .../src/stack-graphs.tsg                      | 59 +++++++++++++------
 .../test/lambdas.py                           |  2 +
 2 files changed, 42 insertions(+), 19 deletions(-)
 create mode 100644 languages/tree-sitter-stack-graphs-python/test/lambdas.py

diff --git a/languages/tree-sitter-stack-graphs-python/src/stack-graphs.tsg b/languages/tree-sitter-stack-graphs-python/src/stack-graphs.tsg
index 48353f14a..f9798fe54 100644
--- a/languages/tree-sitter-stack-graphs-python/src/stack-graphs.tsg
+++ b/languages/tree-sitter-stack-graphs-python/src/stack-graphs.tsg
@@ -759,32 +759,39 @@ inherit .parent_module
 
 (with_clause) {}
 
-(function_definition
-  name: (identifier) @name) {
-  attr (@name.def) node_definition = @name
-}
-
-(function_definition
-  name: (identifier) @name
-  parameters: (parameters) @params
-  body: (block) @body) @func
-{
-  node call
-  node drop_scope
+[
+  (function_definition
+    parameters: (_) @params
+    body: (_) @body
+  ) @func
+  (lambda
+    parameters: (_) @params
+    body: (_) @body
+  )@func
+] {
+  node @func.call
   node return_value
+  node drop_scope
 
-  attr (@name.def) definiens_node = @func
-  edge @func.after_scope -> @name.def
-  edge @name.def -> call
-  edge call -> return_value
+  edge @func.call -> return_value
   edge @body.before_scope -> @params.after_scope
   edge @body.before_scope -> drop_scope
   edge drop_scope -> @func.bottom
   attr (drop_scope) type = "drop_scopes"
-  attr (call) pop_scoped_symbol = "()"
+  attr (@func.call) pop_scoped_symbol = "()"
   edge @params.before_scope -> JUMP_TO_SCOPE_NODE
   attr (return_value) is_exported
   let @func.function_returns = return_value
+}
+
+(function_definition
+  name: (identifier) @name
+  body: (_) @body
+) @func {
+  attr (@name.def) node_definition = @name
+  attr (@name.def) definiens_node = @func
+  edge @func.after_scope -> @name.def
+  edge @name.def -> @func.call
 
   ; Prevent functions defined inside of method bodies from being treated like methods
   let @body.class_self_scope = #null
@@ -802,8 +809,10 @@ inherit .parent_module
   attr (@param.output) push_node = @param
 }
 
-(parameters
-  (_) @param) @params
+[
+  (parameters        (_) @param) @params
+  (lambda_parameters (_) @param) @params
+]
 {
   node @param.param_index
   node @param.param_name
@@ -954,6 +963,18 @@ inherit .parent_module
 ;;
 ;; Expressions
 
+(lambda
+  body: (_) @body
+)@lam {
+  ;; TODO Unify .before_scope, .local_scope, and .input to simplify
+  ;;      uniform treatment of lambdas and function definitions.
+  node @body.before_scope
+  let @body.local_scope = @body.before_scope
+  edge @body.input -> @body.before_scope
+
+  edge @lam.output -> @lam.call
+}
+
 (conditional_expression) {}
 
 (named_expression) {}
diff --git a/languages/tree-sitter-stack-graphs-python/test/lambdas.py b/languages/tree-sitter-stack-graphs-python/test/lambdas.py
new file mode 100644
index 000000000..9812df4de
--- /dev/null
+++ b/languages/tree-sitter-stack-graphs-python/test/lambdas.py
@@ -0,0 +1,2 @@
+sorted([1, 2, 3], key=lambda x: x)
+#                               ^ defined: 1