From ef9f1c8cb2e4f4f17c4f165affe2bdaf434531e5 Mon Sep 17 00:00:00 2001
From: Morgante Pell <morgantep@google.com>
Date: Thu, 18 Jul 2024 11:24:28 -0700
Subject: [PATCH 1/2] add test case

---
 crates/core/src/test.rs | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/crates/core/src/test.rs b/crates/core/src/test.rs
index a76b46d2b..f2dd9051c 100644
--- a/crates/core/src/test.rs
+++ b/crates/core/src/test.rs
@@ -2351,6 +2351,38 @@ fn python_easy_sub() {
     .unwrap();
 }
 
+#[test]
+fn python_type_context() {
+    run_test_expected({
+        TestArgExpected {
+            pattern: r#"
+                |engine marzano(0.1)
+                |language python
+                |
+                |`def matching_method($_) -> $type:
+                | $body` where {
+                |    $type <: `List` => `List[str]`
+                |}
+                |}"#
+            .trim_margin()
+            .unwrap(),
+            source: r#"
+            |def matching_method(bob) -> List:
+            |  print("ok")
+            |"#
+            .trim_margin()
+            .unwrap(),
+            expected: r#"
+            |def matching_method(bob) -> List[str]:
+            |  print("ok")
+            |"#
+            .trim_margin()
+            .unwrap(),
+        }
+    })
+    .unwrap();
+}
+
 #[test]
 fn predicate_maybe() {
     run_test_expected({

From 4c5b0e2cd00353a01c65dea6dc6e1dc133b6d873 Mon Sep 17 00:00:00 2001
From: Morgante Pell <morgantep@google.com>
Date: Thu, 18 Jul 2024 11:39:11 -0700
Subject: [PATCH 2/2] [skip ci] yolo

---
 crates/core/src/test.rs       | 33 ++++++++++++++++++++++++++++++++-
 crates/language/src/python.rs |  1 +
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/crates/core/src/test.rs b/crates/core/src/test.rs
index f2dd9051c..4d75bdc07 100644
--- a/crates/core/src/test.rs
+++ b/crates/core/src/test.rs
@@ -98,6 +98,7 @@ fn match_pattern_libs(
 
     let pattern =
         src_to_problem_libs(pattern, libs, default_language, None, None, None, None)?.problem;
+    println!("PATTERN: {:?}", pattern);
     let results = pattern.execute_file(&RichFile::new(file.to_owned(), src.to_owned()), context);
     let mut execution_result = ExecutionResult {
         input_file_debug_text: "".to_string(),
@@ -2351,6 +2352,37 @@ fn python_easy_sub() {
     .unwrap();
 }
 
+#[test]
+fn tsx_type_context() {
+    run_test_expected({
+        TestArgExpected {
+            pattern: r#"
+                |engine marzano(0.1)
+                |language js
+                |
+                |`function matchingMethod($_): $type {
+                | $body
+                |}` where {
+                |    $type <: `Array` => `Array<string>`
+                |}"#
+            .trim_margin()
+            .unwrap(),
+            source: r#"
+            |function matchingMethod(bob): Array {
+            |  console.log("ok")
+            |}"#
+            .trim_margin()
+            .unwrap(),
+            expected: r#"
+            |function matchingMethod(bob): Array<string> {
+            |  console.log("ok")
+            |}"#
+            .trim_margin()
+            .unwrap(),
+        }
+    })
+    .unwrap();
+}
 #[test]
 fn python_type_context() {
     run_test_expected({
@@ -2362,7 +2394,6 @@ fn python_type_context() {
                 |`def matching_method($_) -> $type:
                 | $body` where {
                 |    $type <: `List` => `List[str]`
-                |}
                 |}"#
             .trim_margin()
             .unwrap(),
diff --git a/crates/language/src/python.rs b/crates/language/src/python.rs
index ff710feb4..968bf5ec0 100644
--- a/crates/language/src/python.rs
+++ b/crates/language/src/python.rs
@@ -66,6 +66,7 @@ impl Language for Python {
 
     fn snippet_context_strings(&self) -> &[(&'static str, &'static str)] {
         &[
+            ("\ndef GRIT_FUNCTION() ->", "\n  pass"),
             ("", ""),
             ("{ ", " }"),
             ("", "\ndef GRIT_FUNCTION():\n    return;"),