From 2b521f036babdc14062b7f8c7e3fb050072046db Mon Sep 17 00:00:00 2001
From: Anthony Sottile <asottile@umich.edu>
Date: Tue, 11 Mar 2025 22:14:52 -0400
Subject: [PATCH] add rudimentary support for PEP 695

---
 grammars/MagicPython.cson            | 138 ++++++++++++++++-
 grammars/MagicPython.tmLanguage      | 223 ++++++++++++++++++++++++++-
 grammars/src/MagicPython.syntax.yaml |  85 +++++++++-
 3 files changed, 440 insertions(+), 6 deletions(-)

diff --git a/grammars/MagicPython.cson b/grammars/MagicPython.cson
index d837d6a..d1f43df 100644
--- a/grammars/MagicPython.cson
+++ b/grammars/MagicPython.cson
@@ -39,6 +39,9 @@ repository:
       {
         include: "#import"
       }
+      {
+        include: "#type-declaration"
+      }
       {
         include: "#class-declaration"
       }
@@ -1099,6 +1102,45 @@ repository:
         ]
       }
     ]
+  "type-declaration":
+    name: "meta.type.python"
+    begin: '''
+      (?x)
+        \\s*(type)\\s+
+          (?=
+            [[:alpha:]_][[:word:]]* \\s* [=\\[]
+          )
+      
+    '''
+    end: "((?=[=]))"
+    beginCaptures:
+      "1":
+        name: "storage.type.type.python"
+    patterns: [
+      {
+        include: "#type-name"
+      }
+      {
+        include: "#type-parameters"
+      }
+    ]
+  "type-name":
+    patterns: [
+      {
+        include: "#illegal-object-name"
+      }
+      {
+        include: "#builtin-possible-callables"
+      }
+      {
+        name: "entity.name.type.type.python"
+        match: '''
+          (?x)
+            \\b ([[:alpha:]_]\\w*) \\b
+          
+        '''
+      }
+    ]
   "class-declaration":
     patterns: [
       {
@@ -1107,7 +1149,7 @@ repository:
           (?x)
             \\s*(class)\\s+
               (?=
-                [[:alpha:]_]\\w* \\s* (:|\\()
+                [[:alpha:]_]\\w* \\s* (:|[\\[(])
               )
           
         '''
@@ -1125,6 +1167,9 @@ repository:
           {
             include: "#class-inheritance"
           }
+          {
+            include: "#type-parameters"
+          }
         ]
       }
     ]
@@ -1355,7 +1400,7 @@ repository:
         \\s*
         (?:\\b(async) \\s+)? \\b(def)\\s+
           (?=
-            [[:alpha:]_][[:word:]]* \\s* \\(
+            [[:alpha:]_][[:word:]]* \\s* [(\\[]
           )
       
     '''
@@ -1372,6 +1417,9 @@ repository:
       {
         include: "#function-def-name"
       }
+      {
+        include: "#type-parameters"
+      }
       {
         include: "#parameters"
       }
@@ -1399,6 +1447,67 @@ repository:
         '''
       }
     ]
+  "type-parameters":
+    name: "meta.function.parameters.type.python"
+    begin: "(\\[)"
+    end: "(\\])"
+    beginCaptures:
+      "1":
+        name: "punctuation.definition.parameters.begin.python"
+    endCaptures:
+      "1":
+        name: "punctuation.definition.parameters.end.python"
+    patterns: [
+      {
+        name: "keyword.operator.unpacking.parameter.python"
+        match: "(\\*\\*|\\*)"
+      }
+      {
+        include: "#lambda-incomplete"
+      }
+      {
+        include: "#illegal-names"
+      }
+      {
+        include: "#illegal-object-name"
+      }
+      {
+        match: '''
+          (?x)
+            ([[:alpha:]_]\\w*)
+              \\s* (?: (,) | (?=[\\]#\\n=]))
+          
+        '''
+        captures:
+          "1":
+            name: "variable.parameter.function.language.python"
+          "2":
+            name: "punctuation.separator.parameters.python"
+      }
+      {
+        include: "#comments"
+      }
+      {
+        include: "#type-loose-default"
+      }
+      {
+        include: "#type-annotated-parameter"
+      }
+    ]
+  "type-loose-default":
+    begin: "(=)"
+    end: "(,)|(?=\\])"
+    beginCaptures:
+      "1":
+        name: "keyword.operator.python"
+    endCaptures:
+      "1":
+        name: "punctuation.separator.parameters.python"
+    patterns: [
+      {
+        include: "#expression"
+      }
+    ]
   parameters:
     name: "meta.function.parameters.python"
     begin: "(\\()"
@@ -1482,6 +1591,31 @@ repository:
         include: "#expression"
       }
     ]
+  "type-annotated-parameter":
+    begin: '''
+      (?x)
+        \\b
+        ([[:alpha:]_]\\w*) \\s* (:)
+      
+    '''
+    end: "(,)|(?=\\])"
+    beginCaptures:
+      "1":
+        name: "variable.parameter.function.language.python"
+      "2":
+        name: "punctuation.separator.annotation.python"
+    endCaptures:
+      "1":
+        name: "punctuation.separator.parameters.python"
+    patterns: [
+      {
+        include: "#expression"
+      }
+      {
+        name: "keyword.operator.assignment.python"
+        match: "=(?!=)"
+      }
+    ]
   "annotated-parameter":
     begin: '''
       (?x)
diff --git a/grammars/MagicPython.tmLanguage b/grammars/MagicPython.tmLanguage
index a1c3f8e..ab1bb16 100644
--- a/grammars/MagicPython.tmLanguage
+++ b/grammars/MagicPython.tmLanguage
@@ -59,6 +59,10 @@
             <key>include</key>
             <string>#import</string>
           </dict>
+          <dict>
+            <key>include</key>
+            <string>#type-declaration</string>
+          </dict>
           <dict>
             <key>include</key>
             <string>#class-declaration</string>
@@ -1675,6 +1679,61 @@ E.g. "arr[idx](args)"
           </dict>
         </array>
       </dict>
+      <key>type-declaration</key>
+      <dict>
+        <key>name</key>
+        <string>meta.type.python</string>
+        <key>begin</key>
+        <string>(?x)
+  \s*(type)\s+
+    (?=
+      [[:alpha:]_][[:word:]]* \s* [=\[]
+    )
+</string>
+        <key>end</key>
+        <string>((?=[=]))</string>
+        <key>beginCaptures</key>
+        <dict>
+          <key>1</key>
+          <dict>
+            <key>name</key>
+            <string>storage.type.type.python</string>
+          </dict>
+        </dict>
+        <key>patterns</key>
+        <array>
+          <dict>
+            <key>include</key>
+            <string>#type-name</string>
+          </dict>
+          <dict>
+            <key>include</key>
+            <string>#type-parameters</string>
+          </dict>
+        </array>
+      </dict>
+      <key>type-name</key>
+      <dict>
+        <key>patterns</key>
+        <array>
+          <dict>
+            <key>include</key>
+            <string>#illegal-object-name</string>
+          </dict>
+          <dict>
+            <key>include</key>
+            <string>#builtin-possible-callables</string>
+          </dict>
+          <dict>
+            <key>name</key>
+            <string>entity.name.type.type.python</string>
+            <key>match</key>
+            <string>(?x)
+  \b ([[:alpha:]_]\w*) \b
+</string>
+          </dict>
+        </array>
+      </dict>
       <key>class-declaration</key>
       <dict>
         <key>patterns</key>
@@ -1686,7 +1745,7 @@ E.g. "arr[idx](args)"
             <string>(?x)
   \s*(class)\s+
     (?=
-      [[:alpha:]_]\w* \s* (:|\()
+      [[:alpha:]_]\w* \s* (:|[\[(])
     )
 </string>
             <key>end</key>
@@ -1717,6 +1776,10 @@ E.g. "arr[idx](args)"
                 <key>include</key>
                 <string>#class-inheritance</string>
               </dict>
+              <dict>
+                <key>include</key>
+                <string>#type-parameters</string>
+              </dict>
             </array>
           </dict>
         </array>
@@ -2103,7 +2166,7 @@ correctly identify the "in" as a control flow keyword.
   \s*
   (?:\b(async) \s+)? \b(def)\s+
     (?=
-      [[:alpha:]_][[:word:]]* \s* \(
+      [[:alpha:]_][[:word:]]* \s* [(\[]
     )
 </string>
         <key>end</key>
@@ -2135,6 +2198,10 @@ correctly identify the "in" as a control flow keyword.
             <key>include</key>
             <string>#function-def-name</string>
           </dict>
+          <dict>
+            <key>include</key>
+            <string>#type-parameters</string>
+          </dict>
           <dict>
             <key>include</key>
             <string>#parameters</string>
@@ -2171,6 +2238,114 @@ correctly identify the "in" as a control flow keyword.
           </dict>
         </array>
       </dict>
+      <key>type-parameters</key>
+      <dict>
+        <key>name</key>
+        <string>meta.function.parameters.type.python</string>
+        <key>begin</key>
+        <string>(\[)</string>
+        <key>end</key>
+        <string>(\])</string>
+        <key>beginCaptures</key>
+        <dict>
+          <key>1</key>
+          <dict>
+            <key>name</key>
+            <string>punctuation.definition.parameters.begin.python</string>
+          </dict>
+        </dict>
+        <key>endCaptures</key>
+        <dict>
+          <key>1</key>
+          <dict>
+            <key>name</key>
+            <string>punctuation.definition.parameters.end.python</string>
+          </dict>
+        </dict>
+        <key>patterns</key>
+        <array>
+          <dict>
+            <key>name</key>
+            <string>keyword.operator.unpacking.parameter.python</string>
+            <key>match</key>
+            <string>(\*\*|\*)</string>
+          </dict>
+          <dict>
+            <key>include</key>
+            <string>#lambda-incomplete</string>
+          </dict>
+          <dict>
+            <key>include</key>
+            <string>#illegal-names</string>
+          </dict>
+          <dict>
+            <key>include</key>
+            <string>#illegal-object-name</string>
+          </dict>
+          <dict>
+            <key>match</key>
+            <string>(?x)
+  ([[:alpha:]_]\w*)
+    \s* (?: (,) | (?=[\]#\n=]))
+</string>
+            <key>captures</key>
+            <dict>
+              <key>1</key>
+              <dict>
+                <key>name</key>
+                <string>variable.parameter.function.language.python</string>
+              </dict>
+              <key>2</key>
+              <dict>
+                <key>name</key>
+                <string>punctuation.separator.parameters.python</string>
+              </dict>
+            </dict>
+          </dict>
+          <dict>
+            <key>include</key>
+            <string>#comments</string>
+          </dict>
+          <dict>
+            <key>include</key>
+            <string>#type-loose-default</string>
+          </dict>
+          <dict>
+            <key>include</key>
+            <string>#type-annotated-parameter</string>
+          </dict>
+        </array>
+      </dict>
+      <key>type-loose-default</key>
+      <dict>
+        <key>begin</key>
+        <string>(=)</string>
+        <key>end</key>
+        <string>(,)|(?=\])</string>
+        <key>beginCaptures</key>
+        <dict>
+          <key>1</key>
+          <dict>
+            <key>name</key>
+            <string>keyword.operator.python</string>
+          </dict>
+        </dict>
+        <key>endCaptures</key>
+        <dict>
+          <key>1</key>
+          <dict>
+            <key>name</key>
+            <string>punctuation.separator.parameters.python</string>
+          </dict>
+        </dict>
+        <key>patterns</key>
+        <array>
+          <dict>
+            <key>include</key>
+            <string>#expression</string>
+          </dict>
+        </array>
+      </dict>
       <key>parameters</key>
       <dict>
         <key>name</key>
@@ -2319,6 +2494,50 @@ correctly identify the "in" as a control flow keyword.
           </dict>
         </array>
       </dict>
+      <key>type-annotated-parameter</key>
+      <dict>
+        <key>begin</key>
+        <string>(?x)
+  \b
+  ([[:alpha:]_]\w*) \s* (:)
+</string>
+        <key>end</key>
+        <string>(,)|(?=\])</string>
+        <key>beginCaptures</key>
+        <dict>
+          <key>1</key>
+          <dict>
+            <key>name</key>
+            <string>variable.parameter.function.language.python</string>
+          </dict>
+          <key>2</key>
+          <dict>
+            <key>name</key>
+            <string>punctuation.separator.annotation.python</string>
+          </dict>
+        </dict>
+        <key>endCaptures</key>
+        <dict>
+          <key>1</key>
+          <dict>
+            <key>name</key>
+            <string>punctuation.separator.parameters.python</string>
+          </dict>
+        </dict>
+        <key>patterns</key>
+        <array>
+          <dict>
+            <key>include</key>
+            <string>#expression</string>
+          </dict>
+          <dict>
+            <key>name</key>
+            <string>keyword.operator.assignment.python</string>
+            <key>match</key>
+            <string>=(?!=)</string>
+          </dict>
+        </array>
+      </dict>
       <key>annotated-parameter</key>
       <dict>
         <key>begin</key>
diff --git a/grammars/src/MagicPython.syntax.yaml b/grammars/src/MagicPython.syntax.yaml
index ebfccda..f6e7ab8 100644
--- a/grammars/src/MagicPython.syntax.yaml
+++ b/grammars/src/MagicPython.syntax.yaml
@@ -142,6 +142,7 @@ repository:
   statement:
     patterns:
       - include: '#import'
+      - include: '#type-declaration'
       - include: '#class-declaration'
       - include: '#function-declaration'
       - include: '#generator'
@@ -831,6 +832,31 @@ repository:
             match: \b(?<!\.)as\b
           - include: '#expression'
 
+  type-declaration:
+    name: meta.type.python
+    begin: |
+      (?x)
+        \s*(type)\s+
+          (?=
+            [[:alpha:]_][[:word:]]* \s* [=\[]
+          )
+
+    end: ((?=[=]))
+    beginCaptures:
+      '1': {name: storage.type.type.python}
+    patterns:
+     - include: '#type-name'
+     - include: '#type-parameters'
+
+  type-name:
+    patterns:
+      - include: '#illegal-object-name'
+      - include: '#builtin-possible-callables'
+      - name: entity.name.type.type.python
+        match: |
+          (?x)
+            \b ([[:alpha:]_]\w*) \b
+
   class-declaration:
     patterns:
       - name: meta.class.python
@@ -838,7 +864,7 @@ repository:
           (?x)
             \s*(class)\s+
               (?=
-                [[:alpha:]_]\w* \s* (:|\()
+                [[:alpha:]_]\w* \s* (:|[\[(])
               )
         end: (:)
         beginCaptures:
@@ -848,6 +874,7 @@ repository:
         patterns:
           - include: '#class-name'
           - include: '#class-inheritance'
+          - include: '#type-parameters'
 
   class-name:
     patterns:
@@ -992,7 +1019,7 @@ repository:
         \s*
         (?:\b(async) \s+)? \b(def)\s+
           (?=
-            [[:alpha:]_][[:word:]]* \s* \(
+            [[:alpha:]_][[:word:]]* \s* [(\[]
           )
 
     end: (:|(?=[#'"\n]))
@@ -1005,6 +1032,7 @@ repository:
 
     patterns:
       - include: '#function-def-name'
+      - include: '#type-parameters'
       - include: '#parameters'
       - include: '#line-continuation'
       - include: '#return-annotation'
@@ -1018,6 +1046,43 @@ repository:
           (?x)
             \b ([[:alpha:]_]\w*) \b
 
+  type-parameters:
+    name: meta.function.parameters.type.python
+    begin: (\[)
+    end: (\])
+    beginCaptures:
+      '1': {name: punctuation.definition.parameters.begin.python}
+    endCaptures:
+      '1': {name: punctuation.definition.parameters.end.python}
+
+    patterns:
+      - name: keyword.operator.unpacking.parameter.python
+        match: (\*\*|\*)
+      - include: '#lambda-incomplete'
+      - include: '#illegal-names'
+      - include: '#illegal-object-name'
+      - match: |
+          (?x)
+            ([[:alpha:]_]\w*)
+              \s* (?: (,) | (?=[\]#\n=]))
+        captures:
+          '1': {name: variable.parameter.function.language.python}
+          '2': {name: punctuation.separator.parameters.python}
+
+      - include: '#comments'
+      - include: '#type-loose-default'
+      - include: '#type-annotated-parameter'
+
+  type-loose-default:
+    begin: (=)
+    end: (,)|(?=\])
+    beginCaptures:
+      '1': {name: keyword.operator.python}
+    endCaptures:
+      '1': {name: punctuation.separator.parameters.python}
+    patterns:
+      - include: '#expression'
+
   parameters:
     name: meta.function.parameters.python
     begin: (\()
@@ -1068,6 +1133,22 @@ repository:
     patterns:
       - include: '#expression'
 
+  type-annotated-parameter:
+    begin: |
+      (?x)
+        \b
+        ([[:alpha:]_]\w*) \s* (:)
+    end: (,)|(?=\])
+    beginCaptures:
+      '1': {name: variable.parameter.function.language.python}
+      '2': {name: punctuation.separator.annotation.python}
+    endCaptures:
+      '1': {name: punctuation.separator.parameters.python}
+    patterns:
+      - include: '#expression'
+      - name: keyword.operator.assignment.python
+        match: =(?!=)
+
   annotated-parameter:
     begin: |
       (?x)