leaf_strategy.py 730 B

12345678910111213141516171819202122232425262728293031323334
  1. from spacy.tokens import Token
  2. from word_processor.types import DEP_TYPES
  3. def leaf_strategy(doc) -> [[Token]]:
  4. """
  5. Should return an arrays of variable names based on leaf strategy
  6. :param doc: spacy document
  7. :return Array of strings
  8. """
  9. for token in doc:
  10. if token.dep_ == DEP_TYPES['ROOT']:
  11. suggestions = dfs(token)
  12. break
  13. return suggestions or []
  14. def dfs(graph, result=None, output=None):
  15. if output is None:
  16. output = []
  17. if result is None:
  18. result = []
  19. flag = False
  20. for u in graph.children:
  21. flag = True
  22. dfs(u, [*result, graph], output)
  23. if flag is False:
  24. output.append([*result, graph])
  25. return output