From 5de5cfdcb8a3e73bb903732f6881758858836b5f Mon Sep 17 00:00:00 2001
From: Yuyao Huang <huangyuyao@outlook.com>
Date: Fri, 19 Apr 2024 13:57:21 +0800
Subject: [PATCH] Expr(); Constant();

---
 tests/test_literals.py                  | 10 ++++++++++
 trace_commentor/commentor.py            | 13 +++++++++++++
 trace_commentor/flags.py                |  2 ++
 trace_commentor/handlers/__init__.py    |  2 ++
 trace_commentor/handlers/definitions.py | 10 +++++++---
 trace_commentor/handlers/expressions.py |  2 ++
 trace_commentor/handlers/literals.py    |  6 ++++++
 trace_commentor/utils.py                | 13 +++++++++++++
 8 files changed, 55 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_literals.py
 create mode 100644 trace_commentor/handlers/expressions.py
 create mode 100644 trace_commentor/handlers/literals.py
 create mode 100644 trace_commentor/utils.py

diff --git a/tests/test_literals.py b/tests/test_literals.py
new file mode 100644
index 0000000..d3ef152
--- /dev/null
+++ b/tests/test_literals.py
@@ -0,0 +1,10 @@
+from trace_commentor import Commentor
+
+
+def test_constant():
+
+    @Commentor()
+    def target():
+        1
+    
+    print(target())
diff --git a/trace_commentor/commentor.py b/trace_commentor/commentor.py
index a1ed5c1..4b968c2 100644
--- a/trace_commentor/commentor.py
+++ b/trace_commentor/commentor.py
@@ -18,6 +18,7 @@ class Commentor(object):
         self._formatters = _formatters + formatters.LIST
         self._lines = []
         self.indent = 0
+        self.state = flags.SOURCE
 
     def __call__(self, func):
 
@@ -60,3 +61,15 @@ class Commentor(object):
             
     def append(self, line):
         self._lines.append(" " * self.indent + str(line))
+            
+    def append_source(self, line):
+        if self.state == flags.COMMENT:
+            self.append('"""')
+        self.append(line)
+        
+    def append_comment(self, line):
+        if self.state == flags.SOURCE:
+            self.append('"""')
+        self.append(line)
+    
+    
diff --git a/trace_commentor/flags.py b/trace_commentor/flags.py
index b82723c..b9fa325 100644
--- a/trace_commentor/flags.py
+++ b/trace_commentor/flags.py
@@ -1,2 +1,4 @@
 DEBUG = True
 INDENT = 4
+SOURCE = 1
+COMMENT = 2
diff --git a/trace_commentor/handlers/__init__.py b/trace_commentor/handlers/__init__.py
index 0164848..141a30d 100644
--- a/trace_commentor/handlers/__init__.py
+++ b/trace_commentor/handlers/__init__.py
@@ -1,2 +1,4 @@
 from .definitions import FunctionDef
 from .statements import Pass
+from .expressions import Expr
+from .literals import Constant
diff --git a/trace_commentor/handlers/definitions.py b/trace_commentor/handlers/definitions.py
index a5b922a..ea58ce8 100644
--- a/trace_commentor/handlers/definitions.py
+++ b/trace_commentor/handlers/definitions.py
@@ -1,9 +1,13 @@
-from ..flags import INDENT
+from .. import flags
+from ..utils import sign, to_source
 
 def FunctionDef(self, cmtor):
-    cmtor.append("def function():")
-    cmtor.indent += INDENT
+    cmtor.append_source(sign("def function():"))
+    cmtor.indent += flags.INDENT
     
     if self is cmtor.root:
         for stmt in self.body:
+            cmtor.append_source(sign(to_source(stmt)))
             cmtor.process(stmt)
+    
+    cmtor.indent -= flags.INDENT
diff --git a/trace_commentor/handlers/expressions.py b/trace_commentor/handlers/expressions.py
new file mode 100644
index 0000000..03b526b
--- /dev/null
+++ b/trace_commentor/handlers/expressions.py
@@ -0,0 +1,2 @@
+def Expr(self, cmtor):
+    cmtor.process(self.value)
diff --git a/trace_commentor/handlers/literals.py b/trace_commentor/handlers/literals.py
new file mode 100644
index 0000000..f53f060
--- /dev/null
+++ b/trace_commentor/handlers/literals.py
@@ -0,0 +1,6 @@
+import ast
+
+from .comments import comment
+
+def Constant(self, cmtor):
+    pass
diff --git a/trace_commentor/utils.py b/trace_commentor/utils.py
new file mode 100644
index 0000000..a3f19f9
--- /dev/null
+++ b/trace_commentor/utils.py
@@ -0,0 +1,13 @@
+import astor
+import inspect
+from . import flags
+
+def sign(line: str):
+    if flags.DEBUG:
+        debug_msg = inspect.currentframe().f_back.f_code.co_name
+        return f"{line}  ---  {debug_msg}"
+    else:
+        return line
+
+def to_source(node):
+    return astor.to_source(node).rstrip("\n")