From 0577c9202c3e7d474ebe27f04bddd9ce237eb75f Mon Sep 17 00:00:00 2001
From: Anthony Kim <anthonykim@microsoft.com>
Date: Mon, 7 Aug 2023 12:30:50 -0700
Subject: [PATCH] receive file content from ts, parse, collect in py

---
 pythonFiles/normalizeSelection.py             | 29 ++++++++++++++++++-
 .../codeExecution/codeExecutionManager.ts     |  3 +-
 src/client/terminals/codeExecution/helper.ts  |  5 ++--
 src/client/terminals/types.ts                 |  2 +-
 4 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/pythonFiles/normalizeSelection.py b/pythonFiles/normalizeSelection.py
index 0363702717ab..e725817add23 100644
--- a/pythonFiles/normalizeSelection.py
+++ b/pythonFiles/normalizeSelection.py
@@ -3,11 +3,18 @@
 
 import ast
 import json
+import os
+import pathlib
 import re
 import sys
 import textwrap
 
+script_dir = pathlib.Path('User/anthonykim/Desktop/vscode-python/pythonFiles/lib/python')
+sys.path.append(os.fspath(script_dir))
+import debugpy
 
+debugpy.connect(5678)
+debugpy.breakpoint()
 def split_lines(source):
     """
     Split selection lines in a version-agnostic way.
@@ -125,6 +132,24 @@ def normalize_lines(selection):
 
     return source
 
+top_level_nodes = [] # collection of top level nodes
+
+class file_node_visitor(ast.NodeVisitor):
+    def visit_nodes(self, node):
+        top_level_nodes.append(node)
+        self.generic_visit(node)
+
+def traverse_file(wholeFileContent):
+    # use ast module to parse content of the file
+    parsed_file_content = ast.parse(wholeFileContent)
+    file_node_visitor().visit(parsed_file_content)
+
+    for node in ast.iter_child_nodes(parsed_file_content):
+        top_level_nodes.append(node)
+        line_start = node.lineno
+        line_end = node.end_lineno
+        code_of_node = ast.get_source
+        # ast.get_source_segment(wholeFileContent, node) This is way to get original code of the selected node
 
 if __name__ == "__main__":
     # Content is being sent from the extension as a JSON object.
@@ -134,7 +159,9 @@ def normalize_lines(selection):
     contents = json.loads(raw.decode("utf-8"))
 
     normalized = normalize_lines(contents["code"])
-
+    normalized_whole_file = normalize_lines(contents["wholeFileContent"])
+    traverse_file(contents["wholeFileContent"]) # traverse file
+    file_node_visitor().visit(ast.parse(contents["wholeFileContent"]))
     # Send the normalized code back to the extension in a JSON object.
     data = json.dumps({"normalized": normalized})
 
diff --git a/src/client/terminals/codeExecution/codeExecutionManager.ts b/src/client/terminals/codeExecution/codeExecutionManager.ts
index 9f1ba6e90d90..5e1f66b89d2d 100644
--- a/src/client/terminals/codeExecution/codeExecutionManager.ts
+++ b/src/client/terminals/codeExecution/codeExecutionManager.ts
@@ -147,7 +147,8 @@ export class CodeExecutionManager implements ICodeExecutionManager {
         }
         const codeExecutionHelper = this.serviceContainer.get<ICodeExecutionHelper>(ICodeExecutionHelper);
         const codeToExecute = await codeExecutionHelper.getSelectedTextToExecute(activeEditor!);
-        const normalizedCode = await codeExecutionHelper.normalizeLines(codeToExecute!);
+        const wholeFileContent = activeEditor!.document.getText();
+        const normalizedCode = await codeExecutionHelper.normalizeLines(codeToExecute!, wholeFileContent);
         if (!normalizedCode || normalizedCode.trim().length === 0) {
             return;
         }
diff --git a/src/client/terminals/codeExecution/helper.ts b/src/client/terminals/codeExecution/helper.ts
index 0d5694b4a28d..68d5aa63e8af 100644
--- a/src/client/terminals/codeExecution/helper.ts
+++ b/src/client/terminals/codeExecution/helper.ts
@@ -33,7 +33,7 @@ export class CodeExecutionHelper implements ICodeExecutionHelper {
         this.interpreterService = serviceContainer.get<IInterpreterService>(IInterpreterService);
     }
 
-    public async normalizeLines(code: string, resource?: Uri): Promise<string> {
+    public async normalizeLines(code: string, wholeFileContent: string, resource?: Uri): Promise<string> {
         try {
             if (code.trim().length === 0) {
                 return '';
@@ -66,7 +66,7 @@ export class CodeExecutionHelper implements ICodeExecutionHelper {
 
             // The normalization script expects a serialized JSON object, with the selection under the "code" key.
             // We're using a JSON object so that we don't have to worry about encoding, or escaping non-ASCII characters.
-            const input = JSON.stringify({ code });
+            const input = JSON.stringify({ code, wholeFileContent });
             observable.proc?.stdin?.write(input);
             observable.proc?.stdin?.end();
 
@@ -110,6 +110,7 @@ export class CodeExecutionHelper implements ICodeExecutionHelper {
 
         const { selection } = textEditor;
         let code: string;
+        const wholeFileContent = textEditor.document.getText(); // This is a way to get the whole text content from the user
         if (selection.isEmpty) {
             code = textEditor.document.lineAt(selection.start.line).text;
         } else if (selection.isSingleLine) {
diff --git a/src/client/terminals/types.ts b/src/client/terminals/types.ts
index 47ac16d9e08b..585fea913166 100644
--- a/src/client/terminals/types.ts
+++ b/src/client/terminals/types.ts
@@ -15,7 +15,7 @@ export interface ICodeExecutionService {
 export const ICodeExecutionHelper = Symbol('ICodeExecutionHelper');
 
 export interface ICodeExecutionHelper {
-    normalizeLines(code: string): Promise<string>;
+    normalizeLines(code: string, wholeFileContent: string): Promise<string>;
     getFileToExecute(): Promise<Uri | undefined>;
     saveFileIfDirty(file: Uri): Promise<Resource>;
     getSelectedTextToExecute(textEditor: TextEditor): Promise<string | undefined>;