Skip to content

Instantly share code, notes, and snippets.

@rjmholt
Created August 11, 2020 16:47
Show Gist options
  • Save rjmholt/9530c9a5d2ca1a31f8c49546ce9159c6 to your computer and use it in GitHub Desktop.
Save rjmholt/9530c9a5d2ca1a31f8c49546ce9159c6 to your computer and use it in GitHub Desktop.
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