PowerShell AST visitor to break up a file by comments
using System; | |
using System.Collections.Generic; | |
using System.Management.Automation.Language; | |
public class ScriptExtent : IScriptExtent | |
{ | |
private readonly IScriptPosition _start; | |
private readonly IScriptPosition _end; | |
private readonly Lazy<string> _textLazy; | |
public ScriptExtent(IScriptPosition start, IScriptPosition end) | |
{ | |
_start = start; | |
_end = end; | |
_textLazy = new Lazy<string>(GetText); | |
} | |
public int EndColumnNumber => _end.ColumnNumber; | |
public int EndLineNumber => _end.LineNumber; | |
public int EndOffset => _end.Offset; | |
public IScriptPosition EndScriptPosition => _end; | |
public string File => _start.File; | |
public int StartColumnNumber => _start.ColumnNumber; | |
public int StartLineNumber => _start.LineNumber; | |
public int StartOffset => _start.Offset; | |
public IScriptPosition StartScriptPosition => _start; | |
public string Text => _textLazy.Value; | |
private string GetText() | |
{ | |
string scriptText = _start.GetFullScript(); | |
return scriptText?.Substring(_start.Offset, _end.Offset - _start.Offset); | |
} | |
} | |
public class CellFindingVisitor : AstVisitor2 | |
{ | |
public static List<IScriptExtent> GetCodeRegionsFromInput(string input) | |
{ | |
Ast ast = Parser.ParseInput(input, out Token[] tokens, out _); | |
return GetCodeRegions(ast, tokens); | |
} | |
public static List<IScriptExtent> GetCodeRegionsFromFile(string filePath) | |
{ | |
Ast ast = Parser.ParseFile(filePath, out Token[] tokens, out _); | |
return GetCodeRegions(ast, tokens); | |
} | |
private static List<IScriptExtent> GetCodeRegions(Ast ast, IReadOnlyList<Token> tokens) | |
{ | |
var visitor = new CellFindingVisitor(ast, tokens); | |
ast.Visit(visitor); | |
visitor.ProcessScriptEnd(); | |
return visitor._cellExtents; | |
} | |
private readonly Ast _ast; | |
private readonly IReadOnlyList<Token> _tokens; | |
private readonly List<IScriptExtent> _cellExtents; | |
private int _tokenIndex; | |
private IScriptPosition _currentCellStart; | |
private CellFindingVisitor(Ast ast, IReadOnlyList<Token> tokens) | |
{ | |
_ast = ast; | |
_tokens = tokens; | |
_tokenIndex = 0; | |
_cellExtents = new List<IScriptExtent>(); | |
_currentCellStart = ast.Extent.StartScriptPosition; | |
} | |
public override AstVisitAction VisitAssignmentStatement(AssignmentStatementAst assignmentStatementAst) | |
=> VisitStatement(assignmentStatementAst); | |
public override AstVisitAction VisitFunctionDefinition(FunctionDefinitionAst functionDefinitionAst) | |
=> VisitStatement(functionDefinitionAst); | |
public override AstVisitAction VisitSwitchStatement(SwitchStatementAst switchStatementAst) | |
=> VisitStatement(switchStatementAst); | |
public override AstVisitAction VisitThrowStatement(ThrowStatementAst throwStatementAst) | |
=> VisitStatement(throwStatementAst); | |
public override AstVisitAction VisitTrap(TrapStatementAst trapStatementAst) | |
=> VisitStatement(trapStatementAst); | |
public override AstVisitAction VisitTryStatement(TryStatementAst tryStatementAst) | |
=> VisitStatement(tryStatementAst); | |
public override AstVisitAction VisitTypeDefinition(TypeDefinitionAst typeDefinitionAst) | |
=> VisitStatement(typeDefinitionAst); | |
public override AstVisitAction VisitUsingStatement(UsingStatementAst usingStatementAst) | |
=> VisitStatement(usingStatementAst); | |
public override AstVisitAction VisitWhileStatement(WhileStatementAst whileStatementAst) | |
=> VisitStatement(whileStatementAst); | |
public override AstVisitAction VisitBlockStatement(BlockStatementAst blockStatementAst) | |
=> VisitStatement(blockStatementAst); | |
public override AstVisitAction VisitBreakStatement(BreakStatementAst breakStatementAst) | |
=> VisitStatement(breakStatementAst); | |
public override AstVisitAction VisitConfigurationDefinition(ConfigurationDefinitionAst configurationDefinitionAst) | |
=> VisitStatement(configurationDefinitionAst); | |
public override AstVisitAction VisitContinueStatement(ContinueStatementAst continueStatementAst) | |
=> VisitStatement(continueStatementAst); | |
public override AstVisitAction VisitDataStatement(DataStatementAst dataStatementAst) | |
=> VisitStatement(dataStatementAst); | |
public override AstVisitAction VisitDoUntilStatement(DoUntilStatementAst doUntilStatementAst) | |
=> VisitStatement(doUntilStatementAst); | |
public override AstVisitAction VisitDoWhileStatement(DoWhileStatementAst doWhileStatementAst) | |
=> VisitStatement(doWhileStatementAst); | |
public override AstVisitAction VisitDynamicKeywordStatement(DynamicKeywordStatementAst dynamicKeywordStatementAst) | |
=> VisitStatement(dynamicKeywordStatementAst); | |
public override AstVisitAction VisitExitStatement(ExitStatementAst exitStatementAst) | |
=> VisitStatement(exitStatementAst); | |
public override AstVisitAction VisitForEachStatement(ForEachStatementAst forEachStatementAst) | |
=> VisitStatement(forEachStatementAst); | |
public override AstVisitAction VisitForStatement(ForStatementAst forStatementAst) | |
=> VisitStatement(forStatementAst); | |
public override AstVisitAction VisitIfStatement(IfStatementAst ifStmtAst) | |
=> VisitStatement(ifStmtAst); | |
public override AstVisitAction VisitPipeline(PipelineAst pipelineAst) | |
=> VisitStatement(pipelineAst); | |
public override AstVisitAction VisitReturnStatement(ReturnStatementAst returnStatementAst) | |
=> VisitStatement(returnStatementAst); | |
private AstVisitAction VisitStatement(StatementAst statementAst) | |
{ | |
ProcessCommentsToPosition(statementAst.Extent.StartScriptPosition); | |
SkipTokensToPosition(statementAst.Extent.EndScriptPosition); | |
return AstVisitAction.SkipChildren; | |
} | |
private void ProcessCommentsToPosition(IScriptPosition position) | |
{ | |
IScriptPosition firstCommentPosition = null; | |
for (; _tokenIndex < _tokens.Count; _tokenIndex++) | |
{ | |
Token currToken = _tokens[_tokenIndex]; | |
if (currToken.Extent.StartOffset >= position.Offset) | |
{ | |
break; | |
} | |
if (currToken.Kind == TokenKind.Comment) | |
{ | |
firstCommentPosition = currToken.Extent.StartScriptPosition; | |
break; | |
} | |
} | |
if (firstCommentPosition != null) | |
{ | |
_cellExtents.Add(new ScriptExtent(_currentCellStart, firstCommentPosition)); | |
_currentCellStart = position; | |
} | |
} | |
private void SkipTokensToPosition(IScriptPosition position) | |
{ | |
for (; _tokenIndex < _tokens.Count; _tokenIndex++) | |
{ | |
Token currToken = _tokens[_tokenIndex]; | |
if (currToken.Extent.StartOffset > position.Offset) | |
{ | |
break; | |
} | |
} | |
} | |
private void ProcessScriptEnd() | |
{ | |
IScriptPosition finalCommentStartPosition = null; | |
for (; _tokenIndex < _tokens.Count; _tokenIndex++) | |
{ | |
Token currToken = _tokens[_tokenIndex]; | |
if (currToken.Kind == TokenKind.Comment) | |
{ | |
finalCommentStartPosition = currToken.Extent.StartScriptPosition; | |
break; | |
} | |
} | |
_cellExtents.Add(new ScriptExtent(_currentCellStart, finalCommentStartPosition ?? _ast.Extent.EndScriptPosition)); | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment