#86 - Adds enclosing ranges to function and class definitions

packages/pyright-scip/snapshots/input/enclosing_range_class/advanced.py

@@ -0,0 +1,13 @@
+# format-options: showRanges
+def class_decorator(cls):
+    def wrapper(*args, **kwargs):
+        return cls(*args, **kwargs)
+    return wrapper
+
+@class_decorator
+class Test:
+    def __init__(self, x: float):
+        self.x = x
+
+    def test(self) -> float:
+        return self.x

packages/pyright-scip/snapshots/input/enclosing_range_class/simple.py

@@ -0,0 +1,3 @@
+# format-options: showRanges
+class Test():
+  pass

packages/pyright-scip/snapshots/input/enclosing_range_function/advanced.py

@@ -0,0 +1,11 @@
+# format-options: showRanges
+def decorator(func):
+    def wrapper(*args, **kwargs):
+        return func(*args, **kwargs)
+    return wrapper
+
+@decorator
+def func(x: float) -> float:
+    test = x
+
+    return test

packages/pyright-scip/snapshots/input/enclosing_range_function/simple.py

@@ -0,0 +1,3 @@
+# format-options: showRanges
+def simple(x):
+  return x

packages/pyright-scip/snapshots/output/enclosing_range_class/advanced.py

@@ -0,0 +1,50 @@
+# < definition scip-python python snapshot-util 0.1 advanced/__init__:
+
+# format-options: showRanges
+# ⌄ start enclosing_range scip-python python snapshot-util 0.1 advanced/class_decorator().
+def class_decorator(cls):
+#   ^^^^^^^^^^^^^^^ definition  snapshot-util 0.1 advanced/class_decorator().
+#                   ^^^ definition  snapshot-util 0.1 advanced/class_decorator().(cls)
+#   ⌄ start enclosing_range scip-python python snapshot-util 0.1 advanced/class_decorator().wrapper().
+    def wrapper(*args, **kwargs):
+#       ^^^^^^^ definition  snapshot-util 0.1 advanced/class_decorator().wrapper().
+#                ^^^^ definition  snapshot-util 0.1 advanced/class_decorator().wrapper().(args)
+#                        ^^^^^^ definition  snapshot-util 0.1 advanced/class_decorator().wrapper().(kwargs)
+        return cls(*args, **kwargs)
+#              ^^^ reference  snapshot-util 0.1 advanced/class_decorator().(cls)
+#                   ^^^^ reference  snapshot-util 0.1 advanced/class_decorator().wrapper().(args)
+#                           ^^^^^^ reference  snapshot-util 0.1 advanced/class_decorator().wrapper().(kwargs)
+#   ⌃ end enclosing_range scip-python python snapshot-util 0.1 advanced/class_decorator().wrapper().
+    return wrapper
+#          ^^^^^^^ reference  snapshot-util 0.1 advanced/class_decorator().wrapper().
+# ⌃ end enclosing_range scip-python python snapshot-util 0.1 advanced/class_decorator().
+
+# ⌄ start enclosing_range scip-python python snapshot-util 0.1 advanced/Test#
+@class_decorator
+#^^^^^^^^^^^^^^^ reference  snapshot-util 0.1 advanced/class_decorator().
+class Test:
+#     ^^^^ definition  snapshot-util 0.1 advanced/Test#
+#     ^^^^ definition  snapshot-util 0.1 advanced/Test#
+#   ⌄ start enclosing_range scip-python python snapshot-util 0.1 advanced/Test#__init__().
+    def __init__(self, x: float):
+#       ^^^^^^^^ definition  snapshot-util 0.1 advanced/Test#__init__().
+#                ^^^^ definition  snapshot-util 0.1 advanced/Test#__init__().(self)
+#                      ^ definition  snapshot-util 0.1 advanced/Test#__init__().(x)
+#                         ^^^^^ reference  python-stdlib 3.11 builtins/float#
+        self.x = x
+#       ^^^^ reference  snapshot-util 0.1 advanced/Test#__init__().(self)
+#            ^ definition  snapshot-util 0.1 advanced/Test#x.
+#                ^ reference  snapshot-util 0.1 advanced/Test#__init__().(x)
+#   ⌃ end enclosing_range scip-python python snapshot-util 0.1 advanced/Test#__init__().
+
+#   ⌄ start enclosing_range scip-python python snapshot-util 0.1 advanced/Test#test().
+    def test(self) -> float:
+#       ^^^^ definition  snapshot-util 0.1 advanced/Test#test().
+#            ^^^^ definition  snapshot-util 0.1 advanced/Test#test().(self)
+#                     ^^^^^ reference  python-stdlib 3.11 builtins/float#
+        return self.x
+#              ^^^^ reference  snapshot-util 0.1 advanced/Test#test().(self)
+#                   ^ reference  snapshot-util 0.1 advanced/Test#x.
+#   ⌃ end enclosing_range scip-python python snapshot-util 0.1 advanced/Test#test().
+# ⌃ end enclosing_range scip-python python snapshot-util 0.1 advanced/Test#
+

packages/pyright-scip/snapshots/output/enclosing_range_class/simple.py

@@ -0,0 +1,10 @@
+# < definition scip-python python snapshot-util 0.1 simple/__init__:
+
+# format-options: showRanges
+# ⌄ start enclosing_range scip-python python snapshot-util 0.1 simple/Test#
+class Test():
+#     ^^^^ definition  snapshot-util 0.1 simple/Test#
+#     ^^^^ definition  snapshot-util 0.1 simple/Test#
+  pass
+# ⌃ end enclosing_range scip-python python snapshot-util 0.1 simple/Test#
+

packages/pyright-scip/snapshots/output/enclosing_range_function/advanced.py

@@ -0,0 +1,37 @@
+# < definition scip-python python snapshot-util 0.1 advanced/__init__:
+
+# format-options: showRanges
+# ⌄ start enclosing_range scip-python python snapshot-util 0.1 advanced/decorator().
+def decorator(func):
+#   ^^^^^^^^^ definition  snapshot-util 0.1 advanced/decorator().
+#             ^^^^ definition  snapshot-util 0.1 advanced/decorator().(func)
+#   ⌄ start enclosing_range scip-python python snapshot-util 0.1 advanced/decorator().wrapper().
+    def wrapper(*args, **kwargs):
+#       ^^^^^^^ definition  snapshot-util 0.1 advanced/decorator().wrapper().
+#                ^^^^ definition  snapshot-util 0.1 advanced/decorator().wrapper().(args)
+#                        ^^^^^^ definition  snapshot-util 0.1 advanced/decorator().wrapper().(kwargs)
+        return func(*args, **kwargs)
+#              ^^^^ reference  snapshot-util 0.1 advanced/decorator().(func)
+#                    ^^^^ reference  snapshot-util 0.1 advanced/decorator().wrapper().(args)
+#                            ^^^^^^ reference  snapshot-util 0.1 advanced/decorator().wrapper().(kwargs)
+#   ⌃ end enclosing_range scip-python python snapshot-util 0.1 advanced/decorator().wrapper().
+    return wrapper
+#          ^^^^^^^ reference  snapshot-util 0.1 advanced/decorator().wrapper().
+# ⌃ end enclosing_range scip-python python snapshot-util 0.1 advanced/decorator().
+
+# ⌄ start enclosing_range scip-python python snapshot-util 0.1 advanced/func().
+@decorator
+#^^^^^^^^^ reference  snapshot-util 0.1 advanced/decorator().
+def func(x: float) -> float:
+#   ^^^^ definition  snapshot-util 0.1 advanced/func().
+#        ^ definition  snapshot-util 0.1 advanced/func().(x)
+#           ^^^^^ reference  python-stdlib 3.11 builtins/float#
+#                     ^^^^^ reference  python-stdlib 3.11 builtins/float#
+    test = x
+#   ^^^^ definition local 0
+#          ^ reference  snapshot-util 0.1 advanced/func().(x)
+
+    return test
+#          ^^^^ reference local 0
+# ⌃ end enclosing_range scip-python python snapshot-util 0.1 advanced/func().
+

packages/pyright-scip/snapshots/output/enclosing_range_function/simple.py

@@ -0,0 +1,11 @@
+# < definition scip-python python snapshot-util 0.1 simple/__init__:
+
+# format-options: showRanges
+# ⌄ start enclosing_range scip-python python snapshot-util 0.1 simple/simple().
+def simple(x):
+#   ^^^^^^ definition  snapshot-util 0.1 simple/simple().
+#          ^ definition  snapshot-util 0.1 simple/simple().(x)
+  return x
+#        ^ reference  snapshot-util 0.1 simple/simple().(x)
+# ⌃ end enclosing_range scip-python python snapshot-util 0.1 simple/simple().
+

packages/pyright-scip/src/lib.ts

@@ -40,19 +40,56 @@ export function formatSnapshot(
     const out: string[] = [];
     const symbolTable = getSymbolTable(doc);
 
+    const formatOptionsPrefix = '# format-options:';
+    const formatOptions = {
+        showDocs: false,
+        showRanges: false,
+    };
+
+    for (let line of input.lines) {
+        if (!line.startsWith(formatOptionsPrefix)) {
+            continue;
+        }
+
+        const options = line.slice(formatOptionsPrefix.length).trim().split(',');
+
+        for (let option of options) {
+            const optionName = option.trim();
+
+            if (!(optionName in formatOptions)) {
+                throw new Error(`Invalid format option: ${optionName}`);
+            }
+
+            formatOptions[optionName as keyof typeof formatOptions] = true;
+        }
+
+        break;
+    }
+
     const externalSymbolTable: Map<string, scip.SymbolInformation> = new Map();
     for (let externalSymbol of externalSymbols) {
         externalSymbolTable.set(externalSymbol.symbol, externalSymbol);
     }
 
+    const enclosingRanges: { range: Range; symbol: string }[] = [];
     const symbolsWithDefinitions: Set<string> = new Set();
+
     for (let occurrence of doc.occurrences) {
         const isDefinition = (occurrence.symbol_roles & scip.SymbolRole.Definition) > 0;
         if (isDefinition) {
             symbolsWithDefinitions.add(occurrence.symbol);
         }
+
+        if (occurrence.enclosing_range.length > 0) {
+            enclosingRanges.push({
+                range: Range.fromLsif(occurrence.enclosing_range),
+                symbol: occurrence.symbol,
+            });
+        }
     }
 
+    enclosingRanges.sort(enclosingRangesByLine);
+
     const emittedDocstrings: Set<string> = new Set();
     const pushDoc = (range: Range, symbol: string, isDefinition: boolean, isStartOfLine: boolean) => {
         // Only emit docstrings once
@@ -75,6 +112,10 @@ export function formatSnapshot(
         }
 
         const pushOneDoc = (docs: string[], external: boolean) => {
+            if (!formatOptions.showDocs) {
+                return;
+            }
+
             for (const documentation of docs) {
                 for (const [idx, line] of documentation.split('\n').entries()) {
                     out.push(prefix);
@@ -127,8 +168,33 @@ export function formatSnapshot(
         out.push('\n');
     };
 
+    const pushEnclosingRange = (
+        enclosingRange: {
+            range: Range;
+            symbol: string;
+        },
+        end: boolean = false
+    ) => {
+        if (!formatOptions.showRanges) {
+            return;
+        }
+
+        out.push(commentSyntax);
+        out.push(' '.repeat(Math.max(1, enclosingRange.range.start.character - 1)));
+        if (end) {
+            out.push('⌃ end ');
+        } else {
+            out.push('⌄ start ');
+        }
+        out.push('enclosing_range ');
+        out.push(enclosingRange.symbol);
+        out.push('\n');
+    };
+
     doc.occurrences.sort(occurrencesByLine);
     let occurrenceIndex = 0;
+    const openEnclosingRanges = [];
+
     for (const [lineNumber, line] of input.lines.entries()) {
         // Write 0,0 items ABOVE the first line.
         //  This is the only case where we would need to do this.
@@ -152,6 +218,26 @@ export function formatSnapshot(
             }
         }
 
+        // Check if any enclosing ranges start on this line
+        for (let rangeIndex = 0; rangeIndex < enclosingRanges.length; rangeIndex++) {
+            const enclosingRange = enclosingRanges[rangeIndex];
+
+            if (enclosingRange.range.start.line == lineNumber) {
+                // Switch the range to the open list
+                enclosingRanges.splice(rangeIndex, 1);
+                openEnclosingRanges.push(enclosingRange);
+
+                // Decrement the counter as an item was removed
+                rangeIndex -= 1;
+
+                pushEnclosingRange(enclosingRange);
+
+                continue;
+            }
+
+            break;
+        }
+
         out.push('');
         out.push(line);
         out.push('\n');
@@ -196,6 +282,21 @@ export function formatSnapshot(
 
             pushDoc(range, occurrence.symbol, isDefinition, isStartOfLine);
         }
+
+        for (let openRangeIndex = openEnclosingRanges.length - 1; openRangeIndex >= 0; openRangeIndex--) {
+            const enclosingRange = openEnclosingRanges[openRangeIndex];
+
+            if (enclosingRange.range.end.line == lineNumber) {
+                // Switch the range to the open list
+                openEnclosingRanges.splice(openRangeIndex, 1);
+
+                pushEnclosingRange(enclosingRange, true);
+
+                continue;
+            }
+
+            break;
+        }
     }
     return out.join('');
 }
@@ -231,3 +332,14 @@ export function diffSnapshot(outputPath: string, obtained: string): void {
 function occurrencesByLine(a: scip.Occurrence, b: scip.Occurrence): number {
     return Range.fromLsif(a.range).compare(Range.fromLsif(b.range));
 }
+
+function enclosingRangesByLine(a: { range: Range; symbol: string }, b: { range: Range; symbol: string }): number {
+    // Return the range that starts first, and if they start at the same line, the one that ends last (enclosing).
+    const rangeCompare = a.range.compare(b.range);
+
+    if (rangeCompare !== 0) {
+        return rangeCompare;
+    }
+
+    return b.range.end.line - a.range.end.line;
+}