expand_selection.py

import numpy as np
from glue.config import menubar_plugin
from glue.utils.qt import pick_item
from glue.core.subset import MaskSubsetState


@menubar_plugin("Expand selection")
def my_plugin(session, data_collection):

    if len(data_collection) == 0:
        print("No datasets")
        return

    if len(data_collection) > 1:

        # Get user to pick a dataset
        data = pick_item(list(data_collection),
                         [x.label for x in data_collection],
                         title='Pick a dataset',
                         label='Pick a dataset')

    else:

        data = data_collection[0]

    if len(data.subsets) == 0:
        print("No subsets to expand")
        return

    if len(data.subsets) > 1:

        # Get user to pick a subset
        subset = pick_item(list(data.subsets),
                           [x.label for x in data.subsets],
                           title='Pick a subset',
                           label='Pick a subset')

    else:

        subset = data.subsets[0]

    # Get user to pick an attribute
    attribute = pick_item(list(data.main_components),
                          [x.label for x in data.main_components],
                          title='Pick an attribute',
                          label='Pick an attribute')

    # Find selected categories
    selected_categories = np.unique(subset[attribute])

    # Find new mask
    mask = np.in1d(data[attribute], selected_categories)

    # Make new subset state
    subset_state = MaskSubsetState(mask, data.pixel_component_ids)

    # Update subset state
    subset.subset_state = subset_state