#!/usr/bin/env python
import coral
from coral.util.search.state import QueryState
from coral.util.search.query import match_elem
import coral.util.search.util as U

class MyState(QueryState):
    def visitor(self):
        self.saveResult()
        return True

prj=coral.loadDocument("testVisit.xml")
target=None
query=None

for entry in prj.contents():
    if entry.name == "Target":
        print "found target:",entry.name
        target = entry
    if entry.oclIsKindOf(coral.lang.UML14.Model) and entry.name[:5] == "Query":
        print "found query:",entry.name
        query=entry

assert target and query

query_set=coral.Set()
[query_set.append(a) for a in query.ownedElement]

q_start=None
t_start=None

for elem in query_set:
    if elem.name == "A":
        q_start=elem
assert q_start

for elem in target.ownedElement:
    if elem.name == "A":
        t_start=elem
assert t_start

c=coral.lang.UML14.Class(name="hello")
assert match_elem(q_start, t_start, U.get_all_parts(query_set), MyState)
assert not match_elem(q_start, c, U.get_all_parts(query_set), MyState)
