From a2c901c38518670579b41a87c991c68b0ffd3797 Mon Sep 17 00:00:00 2001
From: Jarrod Overson <jsoverson@gmail.com>
Date: Wed, 22 Jul 2020 13:05:59 -0400
Subject: [PATCH] made statements() work through .body properties

---
 src/misc/util.ts                    |  4 ++++
 src/refactor-session-chainable.ts   | 12 +++++++-----
 test/refactor-chainable-api.test.ts |  8 ++++++++
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/src/misc/util.ts b/src/misc/util.ts
index b553cbb..cdd2153 100644
--- a/src/misc/util.ts
+++ b/src/misc/util.ts
@@ -50,6 +50,10 @@ export function isNodeWithStatements(input: any): input is NodesWithStatements {
   return 'statements' in input;
 }
 
+export function innerBodyStatements(input: any): Node {
+  return 'body' in input ? input.body : input;
+}
+
 export function isLiteral(
   input: any,
 ): input is
diff --git a/src/refactor-session-chainable.ts b/src/refactor-session-chainable.ts
index 11877ec..9e1a3b7 100644
--- a/src/refactor-session-chainable.ts
+++ b/src/refactor-session-chainable.ts
@@ -901,7 +901,9 @@ export class RefactorSessionChainable {
   }
 
   /**
-   * Returns the selects the statements for the selected nodes. Does nothing for nodes that have no statements property.
+   * Returns the selects the statements for the selected nodes. Note: it will "uplevel" the inner statements of nodes with a `.body` property.
+   *
+   * Does nothing for nodes that have no statements property.
    *
    * @example
    *
@@ -928,9 +930,9 @@ export class RefactorSessionChainable {
    * @public
    */
   statements() {
-    const statements: Statement[] = this.filter(isNodeWithStatements).flatMap(
-      (node: NodesWithStatements) => node.statements,
-    );
+    const statements: Statement[] = this.map(innerBodyStatements)
+      .filter(isNodeWithStatements)
+      .flatMap((node: NodesWithStatements) => node.statements);
     return this.$(statements);
   }
 
@@ -1193,7 +1195,7 @@ export function refactor(input: string | Node, options?: GlobalStateOptions): Re
 */
 import {FunctionDeclaration} from 'shift-ast';
 import {PureFunctionAssessment, PureFunctionAssessmentOptions} from './refactor-plugin-unsafe';
-import {isNodeWithStatements} from './misc/util';
+import {isNodeWithStatements, innerBodyStatements} from './misc/util';
 import {parseScript} from 'shift-parser';
 import {matches} from './misc/query';
 const sorry = function(): PureFunctionAssessment | PureFunctionAssessmentOptions | FunctionDeclaration | undefined {
diff --git a/test/refactor-chainable-api.test.ts b/test/refactor-chainable-api.test.ts
index eb1644b..a849b05 100644
--- a/test/refactor-chainable-api.test.ts
+++ b/test/refactor-chainable-api.test.ts
@@ -77,6 +77,14 @@ describe('chainable interface', () => {
     expect($s(firstFn).nameString()).to.equal('foo');
   });
 
+  it('statements should also get .body.statements', () => {
+    const src = `try { var foo = 1; } catch(e){}`;
+    const $s = refactor(src);
+    const innerStatements = $s($s.statements().first('TryCatchStatement')).statements();
+    expect(innerStatements.length).to.equal(1);
+    expect(innerStatements.first().type).to.equal('VariableDeclarationStatement');
+  });
+
   describe('methods w/o arguments', () => {
     it('.delete() should delete self', () => {
       const src = `idExp;function foo(){}\nfoo();`;