Repository: microsoft/qlib
File: qlib/contrib/data/dataset.py (Python)

Code before:
        n_samples (int): number of samples in the same day
        shuffle (bool): whether shuffle data
        drop_last (bool): whether drop last batch < batch_size
        input_size (int): reshape flatten rows as this input_size (backward compatibility)
    """

    def __init__(
        self,
        handler,
        segments,
        seq_len=60,
        horizon=0,
        num_states=0,
        memory_mode="sample",
        batch_size=-1,
        n_samples=None,
        shuffle=True,
        drop_last=False,
        input_size=None,
        **kwargs,
    ):
        if horizon == 0:
            # Try to guess horizon
            if isinstance(handler, dict):
                handler = init_instance_by_config(handler)
            elif isinstance(handler, str):  # pickled handler
                handler = init_instance_by_config(handler)
            label = handler.data_loader.fields["label"][0][0]
            horizon = guess_horizon(label)

        assert num_states == 0 or horizon > 0, "please specify `horizon` to avoid data leakage"
        assert memory_mode in ["sample", "daily"], "unsupported memory mode"
        assert memory_mode == "sample" or batch_size < 0, "daily memory requires daily sampling (`batch_size < 0`)"
        assert batch_size != 0, "invalid batch size"

        if batch_size > 0 and n_samples is not None:
            warnings.warn("`n_samples` can only be used for daily sampling (`batch_size < 0`)")

        self.seq_len = seq_len
        self.horizon = horizon
        self.num_states = num_states
        self.memory_mode = memory_mode
        self.batch_size = batch_size
        self.n_samples = n_samples
        self.shuffle = shuffle
        self.drop_last = drop_last
        self.input_size = input_size
        self.params = (batch_size, n_samples, drop_last, shuffle)  # for train/eval switch

        super().__init__(handler, segments, **kwargs)


Review comment (suggestion):
Could we write it like this?
`if isinstance(handler, (dict, str))`

Code after:
        n_samples (int): number of samples in the same day
        shuffle (bool): whether shuffle data
        drop_last (bool): whether drop last batch < batch_size
        input_size (int): reshape flatten rows as this input_size (backward compatibility)
    """

    def __init__(
        self,
        handler,
        segments,
        seq_len=60,
        horizon=0,
        num_states=0,
        memory_mode="sample",
        batch_size=-1,
        n_samples=None,
        shuffle=True,
        drop_last=False,
        input_size=None,
        **kwargs,
    ):
        if horizon == 0:
            # Try to guess horizon
            if isinstance(handler, (dict, str)):
                handler = init_instance_by_config(handler)
            assert "label" in getattr(handler.data_loader, "fields", None)
            label = handler.data_loader.fields["label"][0][0]
            horizon = guess_horizon([label])

        assert num_states == 0 or horizon > 0, "please specify `horizon` to avoid data leakage"
        assert memory_mode in ["sample", "daily"], "unsupported memory mode"
        assert memory_mode == "sample" or batch_size < 0, "daily memory requires daily sampling (`batch_size < 0`)"
        assert batch_size != 0, "invalid batch size"

        if batch_size > 0 and n_samples is not None:
            warnings.warn("`n_samples` can only be used for daily sampling (`batch_size < 0`)")

        self.seq_len = seq_len
        self.horizon = horizon
        self.num_states = num_states
        self.memory_mode = memory_mode
        self.batch_size = batch_size
        self.n_samples = n_samples
        self.shuffle = shuffle
        self.drop_last = drop_last
        self.input_size = input_size
        self.params = (batch_size, n_samples, drop_last, shuffle)  # for train/eval switch

        super().__init__(handler, segments, **kwargs)

    def setup_data(self, handler_kwargs: dict = None, **kwargs):