REPL Smart Shift+Enter and Dynamic Smart Cursor

package.json

@@ -544,14 +544,16 @@
                             "pythonSurveyNotification",
                             "pythonPromptNewToolsExt",
                             "pythonTerminalEnvVarActivation",
-                            "pythonTestAdapter"
+                            "pythonTestAdapter",
+                            "pythonREPLSmartSend"
                         ],
                         "enumDescriptions": [
                             "%python.experiments.All.description%",
                             "%python.experiments.pythonSurveyNotification.description%",
                             "%python.experiments.pythonPromptNewToolsExt.description%",
                             "%python.experiments.pythonTerminalEnvVarActivation.description%",
-                            "%python.experiments.pythonTestAdapter.description%"
+                            "%python.experiments.pythonTestAdapter.description%",
+                            "%python.experiments.pythonREPLSmartSend.description%"
                         ]
                     },
                     "scope": "machine",
@@ -567,14 +569,16 @@
                             "pythonSurveyNotification",
                             "pythonPromptNewToolsExt",
                             "pythonTerminalEnvVarActivation",
-                            "pythonTestAdapter"
+                            "pythonTestAdapter",
+                            "pythonREPLSmartSend"
                         ],
                         "enumDescriptions": [
                             "%python.experiments.All.description%",
                             "%python.experiments.pythonSurveyNotification.description%",
                             "%python.experiments.pythonPromptNewToolsExt.description%",
                             "%python.experiments.pythonTerminalEnvVarActivation.description%",
-                            "%python.experiments.pythonTestAdapter.description%"
+                            "%python.experiments.pythonTestAdapter.description%",
+                            "%python.experiments.pythonREPLSmartSend.description%"
                         ]
                     },
                     "scope": "machine",

package.nls.json

@@ -44,6 +44,7 @@
     "python.experiments.pythonPromptNewToolsExt.description": "Denotes the Python Prompt New Tools Extension experiment.",
     "python.experiments.pythonTerminalEnvVarActivation.description": "Enables use of environment variables to activate terminals instead of sending activation commands.",
     "python.experiments.pythonTestAdapter.description": "Denotes the Python Test Adapter experiment.",
+    "python.experiments.pythonREPLSmartSend.description": "Denotes the Python REPL Smart Send experiment.",
     "python.formatting.autopep8Args.description": "Arguments passed in. Each argument is a separate item in the array.",
     "python.formatting.autopep8Args.markdownDeprecationMessage": "This setting will soon be deprecated. Please use the [Autopep8 extension](https://marketplace.visualstudio.com/items?itemName=ms-python.autopep8). <br>Learn more [here](https://aka.ms/AAlgvkb).",
     "python.formatting.autopep8Args.deprecationMessage": "This setting will soon be deprecated. Please use the Autopep8 extension. Learn more here: https://aka.ms/AAlgvkb.",

pythonFiles/normalizeSelection.py

@@ -6,6 +6,7 @@
 import re
 import sys
 import textwrap
+from typing import Iterable
 
 
 def split_lines(source):
@@ -118,6 +119,8 @@ def normalize_lines(selection):
 
         # Insert a newline between each top-level statement, and append a newline to the selection.
         source = "\n".join(statements) + "\n"
+        if selection[-2] == "}":
+            source = source[:-1]
     except Exception:
         # If there's a problem when parsing statements,
         # append a blank line to end the block and send it as-is.
@@ -126,17 +129,159 @@ def normalize_lines(selection):
     return source
 
 
+top_level_nodes = []
+min_key = None
+
+
+def check_exact_exist(top_level_nodes, start_line, end_line):
+    exact_nodes = []
+    for node in top_level_nodes:
+        if node.lineno == start_line and node.end_lineno == end_line:
+            exact_nodes.append(node)
+
+    return exact_nodes
+
+
+def traverse_file(wholeFileContent, start_line, end_line, was_highlighted):
+    """
+    Intended to traverse through a user's given file content and find, collect all appropriate lines
+    that should be sent to the REPL in case of smart selection.
+    This could be exact statement such as just a single line print statement,
+    or a multiline dictionary, or differently styled multi-line list comprehension, etc.
+    Then call the normalize_lines function to normalize our smartly selected code block.
+    """
+
+    parsed_file_content = ast.parse(wholeFileContent)
+    smart_code = ""
+    should_run_top_blocks = []
+
+    # Purpose of this loop is to fetch and collect all the
+    # AST top level nodes, and its node.body as child nodes.
+    # Individual nodes will contain information like
+    # the start line, end line and get source segment information
+    # that will be used to smartly select, and send normalized code.
+    for node in ast.iter_child_nodes(parsed_file_content):
+        top_level_nodes.append(node)
+
+        ast_types_with_nodebody = (
+            ast.Module,
+            ast.Interactive,
+            ast.Expression,
+            ast.FunctionDef,
+            ast.AsyncFunctionDef,
+            ast.ClassDef,
+            ast.For,
+            ast.AsyncFor,
+            ast.While,
+            ast.If,
+            ast.With,
+            ast.AsyncWith,
+            ast.Try,
+            ast.Lambda,
+            ast.IfExp,
+            ast.ExceptHandler,
+        )
+        if isinstance(node, ast_types_with_nodebody) and isinstance(
+            node.body, Iterable
+        ):
+            for child_nodes in node.body:
+                top_level_nodes.append(child_nodes)
+
+    exact_nodes = check_exact_exist(top_level_nodes, start_line, end_line)
+
+    # Just return the exact top level line, if present.
+    if len(exact_nodes) > 0:
+        which_line_next = 0
+        for same_line_node in exact_nodes:
+            should_run_top_blocks.append(same_line_node)
+            smart_code += (
+                f"{ast.get_source_segment(wholeFileContent, same_line_node)}\n"
+            )
+            which_line_next = get_next_block_lineno(should_run_top_blocks)
+        return {
+            "normalized_smart_result": smart_code,
+            "which_line_next": which_line_next,
+        }
+
+    # For each of the nodes in the parsed file content,
+    # add the appropriate source code line(s) to be sent to the REPL, dependent on
+    # user is trying to send and execute single line/statement or multiple with smart selection.
+    for top_node in ast.iter_child_nodes(parsed_file_content):
+        if start_line == top_node.lineno and end_line == top_node.end_lineno:
+            should_run_top_blocks.append(top_node)
+
+            smart_code += f"{ast.get_source_segment(wholeFileContent, top_node)}\n"
+            break  # If we found exact match, don't waste computation in parsing extra nodes.
+        elif start_line >= top_node.lineno and end_line <= top_node.end_lineno:
+            # Case to apply smart selection for multiple line.
+            # This is the case for when we have to add multiple lines that should be included in the smart send.
+            # For example:
+            #    'my_dictionary': {
+            #      'Audi': 'Germany',
+            #      'BMW': 'Germany',
+            #      'Genesis': 'Korea',
+            #     }
+            # with the mouse cursor at 'BMW': 'Germany', should send all of the lines that pertains to my_dictionary.
+
+            should_run_top_blocks.append(top_node)
+
+            smart_code += str(ast.get_source_segment(wholeFileContent, top_node))
+            smart_code += "\n"
+
+    normalized_smart_result = normalize_lines(smart_code)
+    which_line_next = get_next_block_lineno(should_run_top_blocks)
+    return {
+        "normalized_smart_result": normalized_smart_result,
+        "which_line_next": which_line_next,
+    }
+
+
+# Look at the last top block added, find lineno for the next upcoming block,
+# This will be used in calculating lineOffset to move cursor in VS Code.
+def get_next_block_lineno(which_line_next):
+    last_ran_lineno = int(which_line_next[-1].end_lineno)
+    next_lineno = int(which_line_next[-1].end_lineno)
+
+    for reverse_node in top_level_nodes:
+        if reverse_node.lineno > last_ran_lineno:
+            next_lineno = reverse_node.lineno
+            break
+    return next_lineno
+
+
 if __name__ == "__main__":
     # Content is being sent from the extension as a JSON object.
     # Decode the data from the raw bytes.
     stdin = sys.stdin if sys.version_info < (3,) else sys.stdin.buffer
     raw = stdin.read()
     contents = json.loads(raw.decode("utf-8"))
+    # Empty highlight means user has not explicitly selected specific text.
+    empty_Highlight = contents.get("emptyHighlight", False)
 
-    normalized = normalize_lines(contents["code"])
+    # We also get the activeEditor selection start line and end line from the typescript VS Code side.
+    # Remember to add 1 to each of the received since vscode starts line counting from 0 .
+    vscode_start_line = contents["startLine"] + 1
+    vscode_end_line = contents["endLine"] + 1
 
     # Send the normalized code back to the extension in a JSON object.
-    data = json.dumps({"normalized": normalized})
+    data = None
+    which_line_next = 0
+
+    if empty_Highlight and contents.get("smartSendExperimentEnabled"):
+        result = traverse_file(
+            contents["wholeFileContent"],
+            vscode_start_line,
+            vscode_end_line,
+            not empty_Highlight,
+        )
+        normalized = result["normalized_smart_result"]
+        which_line_next = result["which_line_next"]
+        data = json.dumps(
+            {"normalized": normalized, "nextBlockLineno": result["which_line_next"]}
+        )
+    else:
+        normalized = normalize_lines(contents["code"])
+        data = json.dumps({"normalized": normalized})
 
     stdout = sys.stdout if sys.version_info < (3,) else sys.stdout.buffer
     stdout.write(data.encode("utf-8"))