From 8869e675cdb4c9b904ecfe1df57e675babe9d661 Mon Sep 17 00:00:00 2001
From: Ned Palacios <npdoesmc@gmail.com>
Date: Thu, 1 Feb 2024 22:07:12 +0800
Subject: [PATCH] feat: expand python language support

---
 languages/python/python_test.go | 33 +++++++++++++++++++---
 languages/python/symbols.txt    | 49 ++++++++++++++++++++++++---------
 2 files changed, 65 insertions(+), 17 deletions(-)

diff --git a/languages/python/python_test.go b/languages/python/python_test.go
index 7720069..aa0bab3 100644
--- a/languages/python/python_test.go
+++ b/languages/python/python_test.go
@@ -13,14 +13,39 @@ func TestPython(t *testing.T) {
 			Name:     "Simple",
 			FileName: "simple.py",
 			Input: `
+class Square:
+	side = 1
+
+	def __init__(self, side):
+		self.side = side
+
+	def area(self):
+		return self.side * self.side
+
+	def perimeter(self):
+		return 4 * self.side
+
 def main():
 	a = 1
 			`,
 			Expected: `
-(tree [0,0 | 0]-[1,6 | 18]
-	(function  [0,0 | 0]-[1,6 | 18]
-		(tree [0,0 | 0]-[1,6 | 18]
-			(assignment a [1,1 | 13]-[1,2 | 14]))))
+(tree [0,0 | 0]-[13,6 | 185]
+	(class Square Square [0,0 | 0]-[10,22 | 165]
+		(tree [1,1 | 15]-[10,22 | 165]
+			(assignment int side [1,1 | 15]-[1,5 | 19])
+			(function any __init__ [3,1 | 26]-[4,18 | 70]
+				(tree [3,1 | 26]-[4,18 | 70]
+					(variable any self [3,13 | 38]-[3,25 | 50])
+					(variable any side [3,13 | 38]-[3,25 | 50])))
+			(function any area [6,1 | 73]-[7,30 | 119]
+				(tree [6,1 | 73]-[7,30 | 119]
+					(variable any self [6,9 | 81]-[6,15 | 87])))
+			(function any perimeter [9,1 | 122]-[10,22 | 165]
+				(tree [9,1 | 122]-[10,22 | 165]
+					(variable any self [9,14 | 135]-[9,20 | 141])))))
+	(function any main [12,0 | 167]-[13,6 | 185]
+		(tree [12,0 | 167]-[13,6 | 185]
+			(assignment int a [13,1 | 180]-[13,2 | 181]))))
 `,
 		},
 	}
diff --git a/languages/python/symbols.txt b/languages/python/symbols.txt
index 181175b..5a3e487 100644
--- a/languages/python/symbols.txt
+++ b/languages/python/symbols.txt
@@ -1,18 +1,41 @@
 (import_statement
-	name: (_) @import.path) @import
+  name: (_) @import.path) @import
+
+(module [(class_definition
+  name: (identifier) @class.name
+  body: (block [
+  (expression_statement
+    (assignment
+      left: (identifier) @variable.name
+      type: (_)? @variable.return-type
+      right: (_) @variable.content) @variable)
+
+  (function_definition
+    name: (identifier) @method.name
+    parameters: (parameters
+      [
+        (identifier) @parameter.name
+        (typed_parameter
+          ((identifier) @parameter.name)
+          type: (_) @parameter.return-type)
+      ]? @parameter) @parameters
+    body: (block
+      (return_statement (_) @block.content)?) @block) @method
+]) @class.body) @class
 
 (function_definition
-	parameters: (parameters
-    	[
-        	(identifier) @parameter.name
-            (typed_parameter
-            	((identifier) @parameter.name)
-                type: (_) @parameter.return-type)
-        ]? @parameter) @parameters) @function
-  	(return_statement
-  		(_) @block.content)?
+  name: (identifier) @function.name
+  parameters: (parameters
+    [
+      (identifier) @parameter.name
+      (typed_parameter
+        ((identifier) @parameter.name)
+          type: (_) @parameter.return-type)
+    ]? @parameter) @parameters
+  body: (block
+    (return_statement (_) @block.content)?) @block) @function])
 
 (expression_statement
-	(assignment
-    	left: (identifier) @assignment.name
-        right: (identifier) @assignment.content) @assignment)
+  (assignment
+      left: (identifier) @assignment.name
+        right: (_) @assignment.content) @assignment)