lean: correctly parse expressions nested within attributes (#1817)

3 files changed, 61 insertions, 23 deletions

diff --git a/pygments/lexers/theorem.py b/pygments/lexers/theorem.py
index beed2e94..529451ad 100644
--- a/pygments/lexers/theorem.py
+++ b/pygments/lexers/theorem.py
@@ -8,7 +8,9 @@
     :license: BSD, see LICENSE for details.
 """
 
-from pygments.lexer import RegexLexer, default, words
+import re
+
+from pygments.lexer import RegexLexer, default, words, include
 from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
     Number, Punctuation, Generic, Whitespace
 
@@ -399,12 +401,34 @@ class LeanLexer(RegexLexer):
     mimetypes = ['text/x-lean']
 
     tokens = {
-        'root': [
+        'expression': [
             (r'\s+', Text),
             (r'/--', String.Doc, 'docstring'),
             (r'/-', Comment, 'comment'),
             (r'--.*?$', Comment.Single),
             (words((
+                    'forall', 'fun', 'Pi', 'from', 'have', 'show', 'assume', 'suffices',
+                    'let', 'if', 'else', 'then', 'in', 'with', 'calc', 'match',
+                    'do'
+                ), prefix=r'\b', suffix=r'\b'), Keyword),
+            (words(('sorry', 'admit'), prefix=r'\b', suffix=r'\b'), Generic.Error),
+            (words(('Sort', 'Prop', 'Type'), prefix=r'\b', suffix=r'\b'), Keyword.Type),
+            (words((
+                '(', ')', ':', '{', '}', '[', ']', '⟨', '⟩', '‹', '›', '⦃', '⦄', ':=', ',',
+            )), Operator),
+            (r'[A-Za-z_\u03b1-\u03ba\u03bc-\u03fb\u1f00-\u1ffe\u2100-\u214f]'
+             r'[.A-Za-z_\'\u03b1-\u03ba\u03bc-\u03fb\u1f00-\u1ffe\u2070-\u2079'
+             r'\u207f-\u2089\u2090-\u209c\u2100-\u214f0-9]*', Name),
+            (r'0x[A-Za-z0-9]+', Number.Integer),
+            (r'0b[01]+', Number.Integer),
+            (r'\d+', Number.Integer),
+            (r'"', String.Double, 'string'),
+            (r"'(?:(\\[\\\"'nt])|(\\x[0-9a-fA-F]{2})|(\\u[0-9a-fA-F]{4})|.)'", String.Char),
+            (r'[~?][a-z][\w\']*:', Name.Variable),
+            (r'\S', Name.Builtin.Pseudo),
+        ],
+        'root': [
+            (words((
                 'import', 'renaming', 'hiding',
                 'namespace',
                 'local',
@@ -439,31 +463,16 @@ class LeanLexer(RegexLexer):
                 'set_option',
                 'run_cmd',
             ), prefix=r'\b', suffix=r'\b'), Keyword.Declaration),
-            (r'@\[[^\]]*\]', Keyword.Declaration),
-            (words((
-                'forall', 'fun', 'Pi', 'from', 'have', 'show', 'assume', 'suffices',
-                'let', 'if', 'else', 'then', 'in', 'with', 'calc', 'match',
-                'do'
-            ), prefix=r'\b', suffix=r'\b'), Keyword),
-            (words(('sorry', 'admit'), prefix=r'\b', suffix=r'\b'), Generic.Error),
-            (words(('Sort', 'Prop', 'Type'), prefix=r'\b', suffix=r'\b'), Keyword.Type),
+            (r'@\[', Keyword.Declaration, 'attribute'),
             (words((
                 '#eval', '#check', '#reduce', '#exit',
                 '#print', '#help',
             ), suffix=r'\b'), Keyword),
-            (words((
-                '(', ')', ':', '{', '}', '[', ']', '⟨', '⟩', '‹', '›', '⦃', '⦄', ':=', ',',
-            )), Operator),
-            (r'[A-Za-z_\u03b1-\u03ba\u03bc-\u03fb\u1f00-\u1ffe\u2100-\u214f]'
-             r'[.A-Za-z_\'\u03b1-\u03ba\u03bc-\u03fb\u1f00-\u1ffe\u2070-\u2079'
-             r'\u207f-\u2089\u2090-\u209c\u2100-\u214f0-9]*', Name),
-            (r'0x[A-Za-z0-9]+', Number.Integer),
-            (r'0b[01]+', Number.Integer),
-            (r'\d+', Number.Integer),
-            (r'"', String.Double, 'string'),
-            (r"'(?:(\\[\\\"'nt])|(\\x[0-9a-fA-F]{2})|(\\u[0-9a-fA-F]{4})|.)'", String.Char),
-            (r'[~?][a-z][\w\']*:', Name.Variable),
-            (r'\S', Name.Builtin.Pseudo),
+            include('expression')
+        ],
+        'attribute': [
+            (r'\]', Keyword.Declaration, '#pop'),
+            include('expression'),
         ],
         'comment': [
             (r'[^/-]', Comment.Multiline),
diff --git a/tests/examplefiles/lean/test.lean b/tests/examplefiles/lean/test.lean
index 2c8b1358..4b719e13 100644
--- a/tests/examplefiles/lean/test.lean
+++ b/tests/examplefiles/lean/test.lean
@@ -207,3 +207,7 @@ let ⟨m, hm⟩ := @zorn α (≤) h (assume a b c, le_trans) in
 ⟨m, assume a ha, le_antisymm (hm a ha) ha⟩
 
 end zorn
+
+-- other bits of tricky syntax
+@[to_additive "See note [foo]"]
+lemma mul_one : sorry := sorry
diff --git a/tests/examplefiles/lean/test.lean.output b/tests/examplefiles/lean/test.lean.output
index 2a258da8..dd144778 100644
--- a/tests/examplefiles/lean/test.lean.output
+++ b/tests/examplefiles/lean/test.lean.output
@@ -3604,4 +3604,29 @@
 'end'         Keyword.Declaration
 ' '           Text
 'zorn'        Name
+'\n\n'        Text
+
+'-- other bits of tricky syntax' Comment.Single
+'\n'          Text
+
+'@['          Keyword.Declaration
+'to_additive' Name
+' '           Text
+'"'           Literal.String.Double
+'See note [foo]' Literal.String.Double
+'"'           Literal.String.Double
+']'           Keyword.Declaration
+'\n'          Text
+
+'lemma'       Keyword.Declaration
+' '           Text
+'mul_one'     Name
+' '           Text
+':'           Operator
+' '           Text
+'sorry'       Generic.Error
+' '           Text
+':='          Operator
+' '           Text
+'sorry'       Generic.Error
 '\n'          Text