17 

    17 from cubicweb.rqlrewrite import add_types_restriction

    18 def add_types_restriction(schema, rqlst, newroot=None, solutions=None):

    18 

    19     if newroot is None:

    20         assert solutions is None

    21         if hasattr(rqlst, '_types_restr_added'):

    22             return

    23         solutions = rqlst.solutions

    24         newroot = rqlst

    25         rqlst._types_restr_added = True

    26     else:

    27         assert solutions is not None

    28         rqlst = rqlst.stmt

    29     eschema = schema.eschema

    30     allpossibletypes = {}

    31     for solution in solutions:

    32         for varname, etype in solution.iteritems():

    33             if not varname in newroot.defined_vars or eschema(etype).is_final():

    34                 continue

    35             allpossibletypes.setdefault(varname, set()).add(etype)

    36     for varname in sorted(allpossibletypes):

    37         try:

    38             var = newroot.defined_vars[varname]

    39         except KeyError:

    40             continue

    41         stinfo = var.stinfo

    42         if stinfo.get('uidrels'):

    43             continue # eid specified, no need for additional type specification

    44         try:

    45             typerels = rqlst.defined_vars[varname].stinfo.get('typerels')

    46         except KeyError:

    47             assert varname in rqlst.aliases

    48             continue

    49         if newroot is rqlst and typerels:

    50             mytyperel = iter(typerels).next()

    51         else:

    52             for vref in newroot.defined_vars[varname].references():

    53                 rel = vref.relation()

    54                 if rel and rel.is_types_restriction():

    55                     mytyperel = rel

    56                     break

    57             else:

    58                 mytyperel = None

    59         possibletypes = allpossibletypes[varname]

    60         if mytyperel is not None:

    61             # variable as already some types restriction. new possible types

    62             # can only be a subset of existing ones, so only remove no more

    63             # possible types

    64             for cst in mytyperel.get_nodes(Constant):

    65                 if not cst.value in possibletypes:

    66                     cst.parent.remove(cst)

    67                     try:

    68                         stinfo['possibletypes'].remove(cst.value)

    69                     except KeyError:

    70                         # restriction on a type not used by this query, may

    71                         # occurs with X is IN(...)

    72                         pass

    73         else:

    74             # we have to add types restriction

    75             if stinfo.get('scope') is not None:

    76                 rel = var.scope.add_type_restriction(var, possibletypes)

    77             else:

    78                 # tree is not annotated yet, no scope set so add the restriction

    79                 # to the root

    80                 rel = newroot.add_type_restriction(var, possibletypes)

    81             stinfo['typerels'] = frozenset((rel,))

    82             stinfo['possibletypes'] = possibletypes