PowerShell AST Modification

As I’ve mentioned in a few posts the AST is now public in PowerShell 3.0. Each node in the AST derives from the AST class which in turn implements the Visit method. The Visit method accepts an ICustomAstVisitor instance that is used to visit the different nodes within the AST. This gives us a chance to change the make up of the script block. In order to get a hold of an AST to process, we can use the Parser class.

Built In AST Visitor

There is a built in base class that we can utilize to easily implement the an AST visitor. Interestingly, it’s called AstVisitor. It offers a basic way of implementing an AST visitor that removes the need to implement each and every one of the AST visitation methods. Since there is a method for every type of language element found in PowerShell, this can be a pretty handy starting point. Additionally, there is an overload for the AST Visit method that enables this class to easily visit each node in the tree. The only issue with this type of visitation is that it is for reading the AST but not modifying it.

Notice that there is a difference between the VisitScriptBlock method on the AstVisitor base class and the ICustomAstVisitor method. Rather than returning a modified AST with a regular visitor type, we instead have the chance to return an action based on the state of the current script block. It allows us to visit the AST but not modify it. It is a good method of analyzing the script, but not directly affecting what is within the script. This is why the new IntelliSense features within the ISE is so robust.

Implementing ICustomAstVisitor

To get the most use out of the Visitor pattern implemented into the AST object model in the new version of PowerShell, it requires us to implement the ICustomAstVisitor interface ourselves. Rather than simply visiting the AST, we have a chance to manipulate it. There is a lot more overhead involved in this because the interface defines over 50 methods. Each method accepts the particular AST node type and returns an object. We can utilize the fact that the return type is very generic to return a modified AST object and reconstruct the tree on the fly. To do this we need to implement a reconstruction method for each one of the AST node types.

Above is an example of one of these methods. The provided NamedBlockAst instance is broken down into it’s statements. Then we visit each of the statements. This causes a recursive tree traversal. In the end, we will visit each node in the tree. The return value of the Visit method may be a modified node. The Visit method is not defined within the ICustomAstVisitor interface but is a common way of dealing with an AST visitor implementation. The Visit method is merely responsbile for determining the correct AST visitation method to call for a particular node type.

Once we have each of the base “reconstruction” methods defined, we can then extend our basic custom visitor to do something a bit more interesting.

Modifying The AST

The first step is to extend the basic visitor and override one of the visitation methods to do something interesting. The example below will replace any single quoted string with the string ‘chrome’.

Next, we will create some basic tests that visit and utilize our custom visitor to do so. I created a test that will both fail and succeed. If the custom visitor works correctly, the value returned by the script execution will be ‘chrome’ rather than ‘devenv’.

Now, we can run the tests. Notice that the result is as expected.

The Take Away

So what? Who cares that we can visit the AST? Well for one, we can implement awesome functionality such as the IntelliSense in the ISE. The ability to modify the AST is something entirely different. I’m interested to see where we the community and Microsoft can take this. The custom visitor I created for this post is not complete. I will work on filling it out. Once I have something more complete, I will post the class itself.




You can leave a response, or trackback from your own site.

8 Responses to “PowerShell AST Modification”

  1. [...] PowerShell AST Modification (Adam Driscoll) [...]

  2. James Tryand says:

    A superb article. Nice once Adam.

  3. Just wondering… wouldn’t it allow more robust script modifications, like alias expanding? I think it would work perfectly for that, but I could be wrong. ;)

    Anyway – great article. Things “clicked” more about AST/ Intellisense link. Thank you!

    • adamdriscoll says:

      Hey Bartek,

      Thanks! Yep, it would allow for that. I envision some cool ISE add-ons in the near future :)


  4. [...] my last post I talked about the AST and how it was now possible to modify it using the visitor pattern. After [...]

  5. Jim Przybylinski says:

    Why create a Visit(object) method that determines which method to call next? Isn’t that what Ast.Visit(ICustomAstVister) already does? Instead of calling ‘var ast = Visit(statement)’ you could call ‘var ast = statement.Visit(this)’.

    • adamdriscoll says:

      Yep, you’re right. That would be the correct way to do this. No need to have the Visit method.

Leave a Reply

In an effort to prevent automatic filling, you should perform a task displayed below.

six + 2 =