Indentation for match-case statements

FIX: Properly indent match/case statements.

Author: deephbz

src/python.ts

@@ -15,9 +15,11 @@ function innerBody(context: TreeIndentContext) {
       break
     } else if (before.name == "Comment") {
       pos = before.from
-    } else if (before.name == "Body") {
+    } else if (before.name == "Body" || before.name == "MatchBody") {
       if (context.baseIndentFor(before) + context.unit <= lineIndent) found = before
       node = before
+    } else if (before.name == "MatchClause") {
+      node = before
     } else if (before.type.is("Statement")) {
       node = before
     } else {
@@ -40,7 +42,7 @@ function indentBody(context: TreeIndentContext, node: SyntaxNode) {
   // A normally deindenting keyword that appears at a higher
   // indentation than the block should probably be handled by the next
   // level
-  if (/^\s*(else:|elif |except |finally:)/.test(context.textAfter) && context.lineIndent(context.pos, -1) > base)
+  if (/^\s*(else:|elif |except |finally:|case\s+[^=:]+:)/.test(context.textAfter) && context.lineIndent(context.pos, -1) > base)
     return null
   return base + context.unit
 }
@@ -57,9 +59,20 @@ export const pythonLanguage = LRLanguage.define({
           let inner = innerBody(context)
           return indentBody(context, inner || context.node) ?? context.continue()
         },
+
+        MatchBody: context => {
+          let inner = innerBody(context)
+          return indentBody(context, inner || context.node) ?? context.continue()
+        },
+
         IfStatement: cx => /^\s*(else:|elif )/.test(cx.textAfter) ? cx.baseIndent : cx.continue(),
         "ForStatement WhileStatement": cx => /^\s*else:/.test(cx.textAfter) ? cx.baseIndent : cx.continue(),
         TryStatement: cx => /^\s*(except |finally:|else:)/.test(cx.textAfter) ? cx.baseIndent : cx.continue(),
+        MatchStatement: cx => {
+          if (/^\s*case /.test(cx.textAfter)) return cx.baseIndent + cx.unit
+          return cx.continue()
+        },
+
         "TupleExpression ComprehensionExpression ParamList ArgList ParenthesizedExpression": delimitedIndent({closing: ")"}),
         "DictionaryExpression DictionaryComprehensionExpression SetExpression SetComprehensionExpression": delimitedIndent({closing: "}"}),
         "ArrayExpression ArrayComprehensionExpression": delimitedIndent({closing: "]"}),
@@ -69,6 +82,7 @@ export const pythonLanguage = LRLanguage.define({
           return (inner && indentBody(context, inner)) ?? context.continue()
         }
       }),
+
       foldNodeProp.add({
         "ArrayExpression DictionaryExpression SetExpression TupleExpression": foldInside,
         Body: (node, state) => ({from: node.from + 1, to: node.to - (node.to == state.doc.length ? 0 : 1)})
@@ -82,7 +96,8 @@ export const pythonLanguage = LRLanguage.define({
                        "F", "FR", "RF", "R", "U", "B", "BR", "RB"]
     },
     commentTokens: {line: "#"},
-    indentOnInput: /^\s*([\}\]\)]|else:|elif |except |finally:)$/
+    // Indent logic logic are triggered upon below input patterns
+    indentOnInput: /^\s*([\}\]\)]|else:|elif |except |finally:|case\s+[^:]*:?)$/,
   }
 })
 

test/test-indent.ts

@@ -9,7 +9,7 @@ function check(code: string) {
     let state = EditorState.create({doc: code, extensions: [python().language]})
     for (let pos = 0, lines = code.split("\n"), i = 0; i < lines.length; i++) {
       let line = lines[i], indent = /^\s*/.exec(line)![0].length
-      ist(`${getIndentation(state, pos)} (${i + 1})`, `${indent} (${i + 1})`)
+      ist(`indent=${getIndentation(state, pos)} (line-numer=${i + 1})`, `indent=${indent} (line-numer=${i + 1})`)
       pos += line.length + 1
     }
   }
@@ -45,4 +45,163 @@ try:
 except e:
   bar()
 `))
+
+  it("multi-line-block try-except", check(`
+try:
+  foo()
+  fooz()
+except e:
+  bar()
+  barz()
+finally:
+  baz()
+  bazz()
+`))
+
+
+  it("multi-line-nested-block try-except", check(`
+try:
+  foo()
+  fooz()
+  try:
+    inner()
+    inner2()
+  except e2:
+    f3()
+    f4()
+  else:
+    f5()
+    f6()
+  finally:
+    f7()
+    f8()
+except e:
+  bar()
+  barz()
+finally:
+  baz()
+  bazz()
+`))
+
+  it("match-case", check(`
+match x:
+  case 1:
+    foo()
+  case 2:
+    bar()
+  case _:
+    bar()
+`))
+
+  it("match-case-multi-line-block", check(`
+def func():
+  match x:
+    case 1:
+      foo()
+      fooz()
+    case 2:
+      bar()
+      bar()
+      bar()
+      match y:
+        case 3:
+          bar()
+        case 4:
+          bar()
+    case _:
+      bar()
+`))
+
+  it("class-with-decorators", check(`
+@decorator1
+@decorator2(
+  param1,
+  param2
+)
+class MyClass:
+  def method(self):
+    pass
+`))
+
+  it("list-comprehension", check(`
+result = [
+  x * y
+  for x in range(10)
+  for y in range(5)
+  if x > y
+]
+`))
+
+  it("multi-line-expressions", check(`
+result = (
+  very_long_variable_name +
+  another_long_variable *
+  some_computation(
+    arg1,
+    arg2
+  )
+)
+`))
+
+  it("async-function-and-with", check(`
+async def process_data():
+  async with context() as ctx:
+    result = await ctx.fetch(
+      url,
+      timeout=30
+    )
+    return result
+`))
+
+  it("nested-functions", check(`
+def outer():
+  x = 1
+  def inner1():
+    y = 2
+    def inner2():
+      z = 3
+      return x + y + z
+    return inner2()
+  return inner1()
+`))
+
+  it("type-hints-and-annotations", check(`
+def process_data(
+  data: list[str],
+  config: dict[str, Any]
+) -> tuple[int, str]:
+  result: Optional[str] = None
+  if data:
+    result = data[0]
+  return len(data), result
+`))
+
+  it("multi-line-dict-comprehension", check(`
+config = {
+  key: value
+  for key, value in items
+  if is_valid(
+    key,
+    value
+  )
+}
+`))
+
+  it("multi-line-with-comments", check(`
+def process(
+  x: int,  # The input value
+  y: float  # The coefficient
+):
+  # Compute first step
+  result = x * y
+  
+  # Apply additional processing
+  if result > 0:
+    # Positive case
+    return result
+  else:
+    # Negative case
+    return -result
+`))
+
 })