code: def check_categorical_variables_assignment(
      estimator, needs_group=False
  ):
      """
      Checks that transformers that work only with categorical variables, correctly set
      the values for the attributes `variables` and `variables_`.
  
      The first attribute can take a string, an integer, a list of strings or integers or
      None.
  
      The second attribute can take a list of string or integers, but is assigned after
      checking that the variables are of type object or categorical.
  
      For this check to run, the transformer needs the tag 'variables' set to
      'categorical'.
      """
      # toy df
      X, y = test_df(categorical=True)
  
      # cast one variable as category
      X[["cat_var2"]] = X[["cat_var2"]].astype("category")
  
      # input variables to test
      if needs_group is True:
          _input_vars_ls = [["cat_var1", "cat_var2"], None]
      else:
          _input_vars_ls = [
              "cat_var1",
              ["cat_var1"],
              ["cat_var1", "cat_var2"],
              None,
          ]
  
      # the estimator
      transformer = clone(estimator)
  
      for input_vars in _input_vars_ls:
          # set the different input var examples
          transformer.set_params(variables=input_vars)
  
          # before fitting
          if input_vars is not None:
              assert transformer.variables == input_vars
          else:
              assert transformer.variables is None
  
          # fit
          transformer.fit(X, y)
  
          if input_vars is not None:
              assert transformer.variables == input_vars
  
              if isinstance(input_vars, list):
                  assert transformer.variables_ == input_vars
              else:
                  assert transformer.variables_ == [input_vars]
          else:
              assert transformer.variables is None
              assert transformer.variables_ == ["cat_var1", "cat_var2"]
  
      # test raises error if uses passes numerical variable
      transformer.set_params(variables=["var_1", "cat_var1"])
      with pytest.raises(TypeError):
          transformer.fit(X, y)
  
imports: import pytest
  from sklearn import clone
  from tests.estimator_checks.dataframe_for_checks import test_df
label_desc: Use image node_id 2 for calling a global function with example usage: check_categorical_variables_assignment(estimator, needs_group) without return types