From c67c0cc92b47d166f8a40cee4613937ecc1f42d4 Mon Sep 17 00:00:00 2001
From: Shaun Hamilton <shauhami020@gmail.com>
Date: Mon, 11 Dec 2023 18:57:52 +0200
Subject: [PATCH] feat: add getBlock python helper (#104)

---
 lib/__tests__/python.test.ts | 55 ++++++++++++++++++++++++++++++++++++
 lib/index.ts                 | 36 +++++++++++++++++++++++
 2 files changed, 91 insertions(+)

diff --git a/lib/__tests__/python.test.ts b/lib/__tests__/python.test.ts
index d3f5aef..d2dc008 100644
--- a/lib/__tests__/python.test.ts
+++ b/lib/__tests__/python.test.ts
@@ -24,4 +24,59 @@ def c():
       expect(function_parameters).toEqual("d, e");
     }
   });
+
+  it("getBlock", () => {
+    const code = `
+a = 1
+
+if a == 1:
+  a = 2
+  b = 3
+  if b == 3:
+    a = 4
+
+for i in range(10):
+  a = 1
+`;
+    const matches = [
+      python.getBlock(code, "if a == 1"),
+      python.getBlock(code, /if +\w+ *== *\d+/),
+      // eslint-disable-next-line prefer-regex-literals
+      python.getBlock(code, new RegExp("if +\\w+ *== *\\d+")),
+    ];
+    for (const match of matches) {
+      expect(match).not.toBeNull();
+      if (match) {
+        // eslint-disable-next-line camelcase
+        const { block_indentation, block_body, block_condition } = match;
+        expect(block_condition).toEqual("if a == 1");
+        expect(block_indentation).toEqual(0);
+        expect(block_body).toEqual(
+          `  a = 2
+  b = 3
+  if b == 3:
+    a = 4
+`
+        );
+      }
+    }
+
+    const matches2 = [
+      python.getBlock(code, "for i in range(10)"),
+      python.getBlock(code, /for +\w+ +in +range\(\d+\)/),
+      // eslint-disable-next-line prefer-regex-literals
+      python.getBlock(code, new RegExp("for +\\w+ +in +range\\(\\d+\\)")),
+    ];
+    for (const match of matches2) {
+      expect(match).not.toBeNull();
+      if (match) {
+        // eslint-disable-next-line camelcase
+        const { block_indentation, block_body, block_condition } = match;
+        expect(block_condition).toEqual("for i in range(10)");
+        expect(block_indentation).toEqual(0);
+        expect(block_body).toEqual(`  a = 1
+`);
+      }
+    }
+  });
 });
diff --git a/lib/index.ts b/lib/index.ts
index 266682a..de91e42 100644
--- a/lib/index.ts
+++ b/lib/index.ts
@@ -123,6 +123,42 @@ export module python {
 
     return null;
   }
+
+  /**
+   * Gets a Python block of code matching the `blockPattern`
+   * @param code - Code string to search
+   * @param blockPattern - String or regular expression to match on the block condition
+   *
+   * **Note:** A string `blockPattern` will be escaped to prevent special characters from being treated as regular expression syntax.
+   */
+  export function getBlock(code: string, blockPattern: string | RegExp) {
+    const escapedBlockPattern =
+      blockPattern instanceof RegExp
+        ? blockPattern.source
+        : blockPattern.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
+
+    const regex = new RegExp(
+      `\\n(?<block_indentation> *?)(?<block_condition>${escapedBlockPattern})\\s*:\\n(?<block_body>.*?)(?=\\n\\k<block_indentation>[\\w#]|$)`,
+      "s"
+    );
+
+    const matchedCode = regex.exec(code);
+    if (matchedCode) {
+      /* eslint-disable camelcase */
+      const { block_body, block_indentation, block_condition } =
+        matchedCode.groups;
+
+      const blockIndentationSansNewLine = block_indentation.replace(/\n+/g, "");
+      return {
+        block_body,
+        block_condition,
+        block_indentation: blockIndentationSansNewLine.length,
+      };
+      /* eslint-enable camelcase */
+    }
+
+    return null;
+  }
 }
 
 export class CSSHelp {