# Parse a string to create SpaCy Doc object
import en_core_web_sm

text = 'Forging involves the shaping of metal using localized compressive forces.'

nlp = en_core_web_sm.load()
doc = nlp(text)

# Create the plot
import visualise_spacy_tree
png = visualise_spacy_tree.create_png(doc)

# Write it to a file
with open('parse_tree.png', 'wb') as f:
    f.write(png)

# If you're using Jupyter notebook, you can render it inline
from IPython.display import Image, display
display(Image(png))

# Override node attributes to customise the plot
from spacy.tokens import Token
Token.set_extension('plot', default={})  # Create a token underscore extension
for token in doc:
    node_label = '{0} [{1}])'.format(token.orth_, token.i)
    token._.plot['label'] = node_label
    if token.dep_ == 'ROOT':
        token._.plot['color'] = 'green'

'''
You can set any valid GraphViz dot attribute in 'plot'.
See GraphViz docs for reference of possible node attributes:
https://graphviz.gitlab.io/_pages/doc/info/attrs.html
'''