trait ContextAwarePluginAssignmentTrait in Drupal 9

trait ContextAwarePluginAssignmentTrait {

  /**
   * Ensures the t() method is available.
   *
   * @see \Drupal\Core\StringTranslation\StringTranslationTrait
   */
  protected abstract function t($string, array $args = [], array $options = []);

  /**
   * Wraps the context handler.
   *
   * @return \Drupal\Core\Plugin\Context\ContextHandlerInterface
   */
  protected function contextHandler() {
    return \Drupal::service('context.handler');
  }

  /**
   * Builds a form element for assigning a context to a given slot.
   *
   * @param \Drupal\Core\Plugin\ContextAwarePluginInterface $plugin
   *   The context-aware plugin.
   * @param \Drupal\Component\Plugin\Context\ContextInterface[] $contexts
   *   An array of contexts.
   *
   * @return array
   *   A form element for assigning context.
   */
  protected function addContextAssignmentElement(ContextAwarePluginInterface $plugin, array $contexts) {
    $element = [];
    foreach ($plugin
      ->getContextDefinitions() as $context_slot => $definition) {
      $valid_contexts = $this
        ->contextHandler()
        ->getMatchingContexts($contexts, $definition);
      $options = [];
      foreach ($valid_contexts as $context_id => $context) {
        $element['#tree'] = TRUE;
        $options[$context_id] = $context
          ->getContextDefinition()
          ->getLabel();
        $element[$context_slot] = [
          '#type' => 'value',
          '#value' => $context_id,
        ];
      }

      // Show the context selector only if there is more than 1 option to choose
      // from. Also, show if there is a single option but the plugin does not
      // require a context.
      if (count($options) > 1 || count($options) == 1 && !$definition
        ->isRequired()) {
        $assignments = $plugin
          ->getContextMapping();
        $element[$context_slot] = [
          '#title' => $definition
            ->getLabel() ?: $this
            ->t('Select a @context value:', [
            '@context' => $context_slot,
          ]),
          '#type' => 'select',
          '#options' => $options,
          '#required' => $definition
            ->isRequired(),
          '#default_value' => !empty($assignments[$context_slot]) ? $assignments[$context_slot] : '',
          '#description' => $definition
            ->getDescription(),
        ];
        if (!$definition
          ->isRequired()) {
          $element[$context_slot]['#empty_value'] = '';
        }
      }
    }
    return $element;
  }

}