Skip to content

restrict_by.py

Mixin to allow restriction by ancestor or descendant fields.

RestrictByMixin

Bases: BaseMixin

Mixin to allow restriction by ancestor or descendant fields.

Source code in src/spyglass/utils/mixins/restrict_by.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
class RestrictByMixin(BaseMixin):
    """Mixin to allow restriction by ancestor or descendant fields."""

    _banned_search_tables = set()  # Tables to avoid in restrict_by

    def ban_search_table(self, table):
        """Ban table from search in restrict_by."""
        self._banned_search_tables.update(ensure_names(table, force_list=True))

    def unban_search_table(self, table):
        """Unban table from search in restrict_by."""
        self._banned_search_tables.difference_update(
            ensure_names(table, force_list=True)
        )

    def see_banned_tables(self):
        """Print banned tables."""
        self._info_msg(f"Banned tables: {self._banned_search_tables}")

    def restrict_by(
        self,
        restriction: str = True,
        direction: str = "up",
        return_graph: bool = False,
        verbose: bool = False,
        **kwargs,
    ) -> QueryExpression:
        """Restrict self based on up/downstream table.

        If fails to restrict table, the shortest path may not have been correct.
        If there's a different path that should be taken, ban unwanted tables.

        >>> my_table = MyTable() # must be instantced
        >>> my_table.ban_search_table(UnwantedTable1)
        >>> my_table.ban_search_table([UnwantedTable2, UnwantedTable3])
        >>> my_table.unban_search_table(UnwantedTable3)
        >>> my_table.see_banned_tables()
        >>>
        >>> my_table << my_restriction

        Parameters
        ----------
        restriction : str
            Restriction to apply to the some table up/downstream of self.
        direction : str, optional
            Direction to search for valid restriction. Default 'up'.
        return_graph : bool, optional
            If True, return FindKeyGraph object. Default False, returns
            restricted version of present table.
        verbose : bool, optional
            If True, print verbose output. Default False.

        Returns
        -------
        Union[QueryExpression, TableChain]
            Restricted version of present table or TableChain object. If
            return_graph, use all_ft attribute to see all tables in cascade.
        """
        TableChain = self._graph_deps[0]

        if restriction is True:
            return self

        try:
            ret = self.restrict(restriction)  # Save time trying first
            if len(ret) < len(self):
                # If it actually restricts, if not it might by a dict that
                # is not a valid restriction, returned as True
                self._warn_msg("Restriction valid for this table. Using as is.")
                return ret
        except DataJointError:  # need assert_join_compatible return bool
            self._logger.debug("Restriction not valid. Attempting to cascade.")

        if direction == "up":
            parent, child = None, self
        elif direction == "down":
            parent, child = self, None
        else:
            raise ValueError("Direction must be 'up' or 'down'.")

        graph = TableChain(
            parent=parent,
            child=child,
            direction=direction,
            search_restr=restriction,
            banned_tables=list(self._banned_search_tables),
            cascade=True,
            verbose=verbose,
            **kwargs,
        )

        if not graph.found_restr:
            return None

        if return_graph:
            return graph

        ret = self & graph._get_restr(self.full_table_name)
        warn_text = (
            f" after restrict with path: {graph.path_str}\n\t "
            + "See `help(YourTable.restrict_by)`"
        )
        if len(ret) == len(self):
            self._warn_msg("Same length" + warn_text)
        elif len(ret) == 0:
            self._warn_msg("No entries" + warn_text)

        return ret

    def __lshift__(self, restriction) -> QueryExpression:
        """Restriction by upstream operator e.g. ``q1 << q2``.

        Returns
        -------
        QueryExpression
            A restricted copy of the query expression using the nearest upstream
            table for which the restriction is valid.
        """
        return self.restrict_by(restriction, direction="up")

    def __rshift__(self, restriction) -> QueryExpression:
        """Restriction by downstream operator e.g. ``q1 >> q2``.

        Returns
        -------
        QueryExpression
            A restricted copy of the query expression using the nearest upstream
            table for which the restriction is valid.
        """
        return self.restrict_by(restriction, direction="down")

    def restrict_all(
        self,
        restriction: str = True,
        direction: str = "down",
        return_graph: bool = False,
        skip_external: bool = True,
        verbose: bool = False,
    ) -> List:
        RestrGraph = self._graph_deps[1]
        rg = RestrGraph(
            seed_table=self,
            leaves=dict(
                table_name=self.full_table_name,
                restriction=self.restriction or restriction,
            ),
            direction=direction,
            banned_tables=list(self._banned_search_tables),
            cascade=True,
            skip_external=skip_external,
            verbose=verbose,
        )
        if return_graph:
            return rg
        self._logger.info(rg.restr_ft)
        return None

ban_search_table(table)

Ban table from search in restrict_by.

Source code in src/spyglass/utils/mixins/restrict_by.py 
17
18
19
def ban_search_table(self, table):
    """Ban table from search in restrict_by."""
    self._banned_search_tables.update(ensure_names(table, force_list=True))

unban_search_table(table)

Unban table from search in restrict_by.

Source code in src/spyglass/utils/mixins/restrict_by.py 
21
22
23
24
25
def unban_search_table(self, table):
    """Unban table from search in restrict_by."""
    self._banned_search_tables.difference_update(
        ensure_names(table, force_list=True)
    )

see_banned_tables()

Print banned tables.

Source code in src/spyglass/utils/mixins/restrict_by.py 
27
28
29
def see_banned_tables(self):
    """Print banned tables."""
    self._info_msg(f"Banned tables: {self._banned_search_tables}")

restrict_by(restriction=True, direction='up', return_graph=False, verbose=False, **kwargs)

Restrict self based on up/downstream table.

If fails to restrict table, the shortest path may not have been correct. If there's a different path that should be taken, ban unwanted tables.

my_table = MyTable() # must be instantced my_table.ban_search_table(UnwantedTable1) my_table.ban_search_table([UnwantedTable2, UnwantedTable3]) my_table.unban_search_table(UnwantedTable3) my_table.see_banned_tables()

my_table << my_restriction

Parameters:

Name Type Description Default
restriction str

Restriction to apply to the some table up/downstream of self.

 True
direction str

Direction to search for valid restriction. Default 'up'.

 'up'
return_graph bool

If True, return FindKeyGraph object. Default False, returns restricted version of present table.

 False
verbose bool

If True, print verbose output. Default False.

 False

Returns:

Type Description
Union[QueryExpression, TableChain]

Restricted version of present table or TableChain object. If return_graph, use all_ft attribute to see all tables in cascade.
Source code in src/spyglass/utils/mixins/restrict_by.py 
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
def restrict_by(
    self,
    restriction: str = True,
    direction: str = "up",
    return_graph: bool = False,
    verbose: bool = False,
    **kwargs,
) -> QueryExpression:
    """Restrict self based on up/downstream table.

    If fails to restrict table, the shortest path may not have been correct.
    If there's a different path that should be taken, ban unwanted tables.

    >>> my_table = MyTable() # must be instantced
    >>> my_table.ban_search_table(UnwantedTable1)
    >>> my_table.ban_search_table([UnwantedTable2, UnwantedTable3])
    >>> my_table.unban_search_table(UnwantedTable3)
    >>> my_table.see_banned_tables()
    >>>
    >>> my_table << my_restriction

    Parameters
    ----------
    restriction : str
        Restriction to apply to the some table up/downstream of self.
    direction : str, optional
        Direction to search for valid restriction. Default 'up'.
    return_graph : bool, optional
        If True, return FindKeyGraph object. Default False, returns
        restricted version of present table.
    verbose : bool, optional
        If True, print verbose output. Default False.

    Returns
    -------
    Union[QueryExpression, TableChain]
        Restricted version of present table or TableChain object. If
        return_graph, use all_ft attribute to see all tables in cascade.
    """
    TableChain = self._graph_deps[0]

    if restriction is True:
        return self

    try:
        ret = self.restrict(restriction)  # Save time trying first
        if len(ret) < len(self):
            # If it actually restricts, if not it might by a dict that
            # is not a valid restriction, returned as True
            self._warn_msg("Restriction valid for this table. Using as is.")
            return ret
    except DataJointError:  # need assert_join_compatible return bool
        self._logger.debug("Restriction not valid. Attempting to cascade.")

    if direction == "up":
        parent, child = None, self
    elif direction == "down":
        parent, child = self, None
    else:
        raise ValueError("Direction must be 'up' or 'down'.")

    graph = TableChain(
        parent=parent,
        child=child,
        direction=direction,
        search_restr=restriction,
        banned_tables=list(self._banned_search_tables),
        cascade=True,
        verbose=verbose,
        **kwargs,
    )

    if not graph.found_restr:
        return None

    if return_graph:
        return graph

    ret = self & graph._get_restr(self.full_table_name)
    warn_text = (
        f" after restrict with path: {graph.path_str}\n\t "
        + "See `help(YourTable.restrict_by)`"
    )
    if len(ret) == len(self):
        self._warn_msg("Same length" + warn_text)
    elif len(ret) == 0:
        self._warn_msg("No entries" + warn_text)

    return ret

__lshift__(restriction)

Restriction by upstream operator e.g. q1 << q2.

Returns:

Type Description
QueryExpression

A restricted copy of the query expression using the nearest upstream table for which the restriction is valid.
Source code in src/spyglass/utils/mixins/restrict_by.py 
121
122
123
124
125
126
127
128
129
130
def __lshift__(self, restriction) -> QueryExpression:
    """Restriction by upstream operator e.g. ``q1 << q2``.

    Returns
    -------
    QueryExpression
        A restricted copy of the query expression using the nearest upstream
        table for which the restriction is valid.
    """
    return self.restrict_by(restriction, direction="up")

__rshift__(restriction)

Restriction by downstream operator e.g. q1 >> q2.

Returns:

Type Description
QueryExpression

A restricted copy of the query expression using the nearest upstream table for which the restriction is valid.
Source code in src/spyglass/utils/mixins/restrict_by.py 
132
133
134
135
136
137
138
139
140
141
def __rshift__(self, restriction) -> QueryExpression:
    """Restriction by downstream operator e.g. ``q1 >> q2``.

    Returns
    -------
    QueryExpression
        A restricted copy of the query expression using the nearest upstream
        table for which the restriction is valid.
    """
    return self.restrict_by(restriction, direction="down")