# Copyright (c) 2018-2021, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#
__all__ = ["SelectionGetter"]

from .abstract_batch_position_getter import AbstractBatchPositionGetter
from typing import Any
import weakref
from .graph_model import GraphModel

class SelectionGetter(AbstractBatchPositionGetter):
    """
    Helper to get the selection of the given model. It's supposed to be
    used with GraphModelBatchPosition.
    """

    def __init__(self, model: GraphModel):
        super().__init__(model)

    def __call__(self, drive_item: Any):
        model = self.model
        if model:
            return model.selection
