From 952fa0a15262518047a40dea60642497612a2695 Mon Sep 17 00:00:00 2001
From: Josh Goebel <me@joshgoebel.com>
Date: Sat, 8 May 2021 08:40:01 -0400
Subject: [PATCH] (enh) simplify Python grammar

---
 src/languages/python.js                       | 41 ++++++++++++-------
 test/markup/python-repl/sample.expect.txt     |  2 +-
 test/markup/python/decorators.expect.txt      | 16 ++++----
 .../function-header-comments.expect.txt       |  8 ++--
 test/markup/python/function-header.expect.txt |  2 +-
 test/markup/python/keywords.expect.txt        |  4 +-
 .../python/matrix-multiplication.expect.txt   |  4 +-
 7 files changed, 44 insertions(+), 33 deletions(-)

diff --git a/src/languages/python.js b/src/languages/python.js
index ea765f4651..5ea0dde6bc 100644
--- a/src/languages/python.js
+++ b/src/languages/python.js
@@ -5,6 +5,7 @@ Website: https://www.python.org
 Category: common
 */
 
+import { IDENT_RE } from '../lib/modes.js';
 import * as regex from '../lib/regex.js';
 
 export default function(hljs) {
@@ -375,28 +376,38 @@ export default function(hljs) {
       STRING,
       COMMENT_TYPE,
       hljs.HASH_COMMENT_MODE,
+      {
+        match: [
+          /def/, /\s+/,
+          IDENT_RE
+        ],
+        scope: {
+          1: "keyword",
+          3: "title.function"
+        },
+        contains: [ PARAMS ]
+      },
       {
         variants: [
           {
-            className: 'function',
-            beginKeywords: 'def'
+            match: [
+              /class/, /\s+/,
+              IDENT_RE, /\s*/,
+              /\(\s*/, IDENT_RE,/\s*\)/
+            ],
           },
           {
-            className: 'class',
-            beginKeywords: 'class'
+            match: [
+              /class/, /\s+/,
+              IDENT_RE
+            ],
           }
         ],
-        end: /:/,
-        illegal: /[${=;\n,]/,
-        contains: [
-          hljs.UNDERSCORE_TITLE_MODE,
-          PARAMS,
-          {
-            begin: /->/,
-            endsWithParent: true,
-            keywords: KEYWORDS
-          }
-        ]
+        scope: {
+          1: "keyword",
+          3: "title.class",
+          6: "title.class.inherited",
+        }
       },
       {
         className: 'meta',
diff --git a/test/markup/python-repl/sample.expect.txt b/test/markup/python-repl/sample.expect.txt
index 1f3c7733eb..95dd4c78ad 100644
--- a/test/markup/python-repl/sample.expect.txt
+++ b/test/markup/python-repl/sample.expect.txt
@@ -11,7 +11,7 @@ foo = 42&quot;
 <span class="hljs-meta">&gt;&gt;&gt;</span> <span class="language-python"><span class="hljs-string">&quot;&quot;&quot;</span></span>
 <span class="hljs-meta">...</span> <span class="language-python"><span class="hljs-string">abc</span></span>
 <span class="hljs-meta">...</span> <span class="language-python"><span class="hljs-string">&quot;&quot;&quot;</span></span>
-<span class="hljs-meta">&gt;&gt;&gt;</span> <span class="language-python"><span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">test</span>():</span></span>
+<span class="hljs-meta">&gt;&gt;&gt;</span> <span class="language-python"><span class="hljs-keyword">def</span> <span class="hljs-title function_">test</span>():</span>
 <span class="hljs-meta">...</span> <span class="language-python">    <span class="hljs-keyword">pass</span></span>
 <span class="hljs-meta">&gt;&gt;&gt;</span>
 <span class="hljs-meta">&gt;&gt;&gt;</span>
diff --git a/test/markup/python/decorators.expect.txt b/test/markup/python/decorators.expect.txt
index a0ab79766f..174834925f 100644
--- a/test/markup/python/decorators.expect.txt
+++ b/test/markup/python/decorators.expect.txt
@@ -1,31 +1,31 @@
 <span class="hljs-meta">@foo</span>
-<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">bar</span>():</span>
+<span class="hljs-keyword">def</span> <span class="hljs-title function_">bar</span>():
     <span class="hljs-keyword">pass</span>
 
 <span class="hljs-meta">@foo  </span><span class="hljs-comment"># bar</span>
-<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">baz</span>():</span>
+<span class="hljs-keyword">def</span> <span class="hljs-title function_">baz</span>():
     <span class="hljs-keyword">pass</span>
 
 <span class="hljs-meta">@foo.bar.baz</span>
-<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">qux</span>():</span>
+<span class="hljs-keyword">def</span> <span class="hljs-title function_">qux</span>():
     <span class="hljs-keyword">pass</span>
 
 <span class="hljs-meta">@surround_with(<span class="hljs-params"><span class="hljs-string">&quot;#&quot;</span>, repeat=<span class="hljs-number">3</span></span>)</span>
-<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">text</span>():</span>
+<span class="hljs-keyword">def</span> <span class="hljs-title function_">text</span>():
     <span class="hljs-keyword">return</span> <span class="hljs-string">&quot;hi!&quot;</span>
 
 <span class="hljs-meta">@py38.style</span>
-<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">func</span>():</span>
+<span class="hljs-keyword">def</span> <span class="hljs-title function_">func</span>():
     <span class="hljs-keyword">pass</span>
 
 <span class="hljs-meta">@py[<span class="hljs-string">&quot;3.9&quot;</span>].style</span>
-<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">func</span>():</span>
+<span class="hljs-keyword">def</span> <span class="hljs-title function_">func</span>():
     <span class="hljs-keyword">pass</span>
 
 <span class="hljs-meta">@py[<span class="hljs-number">3.9</span>].style</span>
-<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">func</span>():</span>
+<span class="hljs-keyword">def</span> <span class="hljs-title function_">func</span>():
     <span class="hljs-keyword">pass</span>
 
 <span class="hljs-meta">@<span class="hljs-number">2</span> + <span class="hljs-number">2</span> == <span class="hljs-number">5</span></span>
-<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">func</span>():</span>
+<span class="hljs-keyword">def</span> <span class="hljs-title function_">func</span>():
     <span class="hljs-keyword">pass</span>
diff --git a/test/markup/python/function-header-comments.expect.txt b/test/markup/python/function-header-comments.expect.txt
index 061a1a4742..35041289d8 100644
--- a/test/markup/python/function-header-comments.expect.txt
+++ b/test/markup/python/function-header-comments.expect.txt
@@ -1,10 +1,10 @@
-<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">foo</span>(<span class="hljs-params">
+<span class="hljs-keyword">def</span> <span class="hljs-title function_">foo</span>(<span class="hljs-params">
   bar, <span class="hljs-comment"># commment</span>
-</span>):</span>
+</span>):
     <span class="hljs-keyword">pass</span>
 
 
-<span class="hljs-class"><span class="hljs-keyword">class</span> <span class="hljs-title">Foo</span>(<span class="hljs-params">collections.namedtuple(<span class="hljs-params"><span class="hljs-string">&#x27;Test&#x27;</span></span>), (<span class="hljs-params">
+<span class="hljs-keyword">class</span> <span class="hljs-title class_">Foo</span>(collections.namedtuple(<span class="hljs-string">&#x27;Test&#x27;</span>), (
     <span class="hljs-string">&#x27;name&#x27;</span>, <span class="hljs-comment"># comment</span>
-</span>)</span>):</span>
+)):
     <span class="hljs-keyword">pass</span>
diff --git a/test/markup/python/function-header.expect.txt b/test/markup/python/function-header.expect.txt
index 197b91d7cd..327f2b6118 100644
--- a/test/markup/python/function-header.expect.txt
+++ b/test/markup/python/function-header.expect.txt
@@ -1,2 +1,2 @@
-<span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">f</span>(<span class="hljs-params">x: <span class="hljs-built_in">int</span>, *, y: <span class="hljs-built_in">bool</span> = <span class="hljs-literal">True</span></span>) -&gt; <span class="hljs-literal">None</span>:</span>
+<span class="hljs-keyword">def</span> <span class="hljs-title function_">f</span>(<span class="hljs-params">x: <span class="hljs-built_in">int</span>, *, y: <span class="hljs-built_in">bool</span> = <span class="hljs-literal">True</span></span>) -&gt; <span class="hljs-literal">None</span>:
     <span class="hljs-keyword">pass</span>
diff --git a/test/markup/python/keywords.expect.txt b/test/markup/python/keywords.expect.txt
index 96aa34ebff..501f670464 100644
--- a/test/markup/python/keywords.expect.txt
+++ b/test/markup/python/keywords.expect.txt
@@ -1,5 +1,5 @@
-<span class="hljs-class"><span class="hljs-keyword">class</span> <span class="hljs-title">Shorty</span>(<span class="hljs-params"><span class="hljs-built_in">dict</span></span>):</span>
-    <span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">len</span>(<span class="hljs-params">self</span>):</span>
+<span class="hljs-keyword">class</span> <span class="hljs-title class_">Shorty</span>(<span class="hljs-title class_ inherited__">dict</span>):
+    <span class="hljs-keyword">def</span> <span class="hljs-title function_">len</span>(<span class="hljs-params">self</span>):
         <span class="hljs-keyword">return</span> <span class="hljs-literal">NotImplemented</span>
 
 x = Shorty()
diff --git a/test/markup/python/matrix-multiplication.expect.txt b/test/markup/python/matrix-multiplication.expect.txt
index 256895b24e..91ad4b9261 100644
--- a/test/markup/python/matrix-multiplication.expect.txt
+++ b/test/markup/python/matrix-multiplication.expect.txt
@@ -1,7 +1,7 @@
 <span class="hljs-meta">@meta</span>
-<span class="hljs-class"><span class="hljs-keyword">class</span> <span class="hljs-title">C</span>:</span>
+<span class="hljs-keyword">class</span> <span class="hljs-title class_">C</span>:
 
 <span class="hljs-meta">    @decorator</span>
-    <span class="hljs-function"><span class="hljs-keyword">def</span> <span class="hljs-title">f</span>(<span class="hljs-params">self, H, V, beta, r</span>):</span>
+    <span class="hljs-keyword">def</span> <span class="hljs-title function_">f</span>(<span class="hljs-params">self, H, V, beta, r</span>):
         S = (H @ beta - r).T @ inv(H @ V @ H.T) @ (H @ beta - r)
         <span class="hljs-keyword">return</span> S