diff --git a/src/Context/ContextDefinition.php b/src/Context/ContextDefinition.php
index 722eb8d0..7705b91f 100644
--- a/src/Context/ContextDefinition.php
+++ b/src/Context/ContextDefinition.php
@@ -140,11 +140,10 @@ class ContextDefinition extends ContextDefinitionCore implements ContextDefiniti
       ],
     ];

+    // Find the widget id for this data type.
     foreach ($widgets as $widget_id => $data_types) {
-      foreach ($data_types as $data_type) {
-        if ($dataType == $data_type) {
-          return $widget_id;
-        }
+      if (in_array($dataType, $data_types)) {
+        return $widget_id;
       }
     }

diff --git a/src/Form/ContextFormTrait.php b/src/Form/ContextFormTrait.php
index 51d6e7e2..d27057ea 100644
--- a/src/Form/ContextFormTrait.php
+++ b/src/Form/ContextFormTrait.php
@@ -57,8 +57,10 @@ trait ContextFormTrait {
     }

     $form['context'][$context_name]['setting'] = [
-      '#type' => 'textarea',
       '#title' => $title,
+      // At least one test should fail when #required is not set here, as the
+      // new widget has a default of required. No tests fail without it though.
+      // @todo Add test to check #required value is working correctly.
       '#required' => $context_definition->isRequired(),
       '#default_value' => $default_value,
     ];
@@ -70,18 +72,21 @@ trait ContextFormTrait {
       $typed_data = $dataManager->create($dataDefinition);
 
       if (isset($configuration['context_mapping']['data']) && $property_path = $configuration['context_mapping']['data']) {
-        $sub_paths = explode('.', $property_path);
-        $dataType = \Drupal::entityManager()->getStorage('field_storage_config')->load($sub_paths[0] . '.' . $sub_paths[1])->getType();
+        list($entity_type, $field_name) = explode('.', $property_path);
+        $dataType = \Drupal::service('entity_field.manager')->getFieldStorageDefinitions($entity_type)[$field_name]->getType();

         if ($widget_id = $context_definition->getWidgetId($dataType)) {
           $widget = $this->getFormWidgetManager()->createInstance($widget_id);
           $sub_form = [];
           $sub_form_state = SubformState::createForSubform($sub_form, $form, $form_state);

-          $form['context'][$context_name]['setting'] = $widget->form($typed_data, $sub_form_state);
+          // Add the widget form to the existing $context_name setting form.
+          $form['context'][$context_name]['setting'] += $widget->form($typed_data, $sub_form_state)['value'];
         }
       }
     }
+    // Set default type to simple 'textfield' if no widget was matched above.
+    $form['context'][$context_name]['setting'] += ['#type' => 'textfield'];

     $element = &$form['context'][$context_name]['setting'];