diff --git a/token.module b/token.module
index 62a06db..22a938e 100644
--- a/token.module
+++ b/token.module
@@ -10,11 +10,13 @@ use Drupal\Component\Utility\Unicode;
 use Drupal\Component\Utility\Xss;
 use Drupal\Core\Block\BlockPluginInterface;
 use Drupal\Core\Cache\Cache;
+use Drupal\Core\Entity\Display\EntityViewDisplayInterface;
 use Drupal\Core\Form\FormStateInterface;
 use Drupal\Core\Render\Element;
 use Drupal\Core\Routing\RouteMatchInterface;
 use Drupal\Core\Utility\Token;
 use Drupal\Core\Url;
+use Drupal\field\Entity\FieldStorageConfig;
 
 /**
  * The maximum depth for token tree recursion.
@@ -195,26 +197,40 @@ function token_field_info_alter(&$info) {
 }
 
 /**
- * Implements hook_field_display_alter().
+ * Implements hook_entity_view_display_alter().
  */
-function token_field_display_alter(&$display, $context) {
+function token_entity_view_display_alter(EntityViewDisplayInterface $entity_display, array $context) {
   if ($context['view_mode'] == 'token') {
-    $view_display = entity_get_display($context['instance']['entity_type'], $context['instance']['bundle'], $context['view_mode']);
+    /** @var EntityViewDisplayInterface $view_display */
+    $view_display = \Drupal::entityManager()->getStorage('entity_view_display')
+      ->load($context['entity_type'] . '.' . $context['bundle'] . '.' . $context['view_mode']);
+
     // If the token view mode fell back to the 'default' view mode, then
     // use the default token formatter.
     if (!$view_display->status()) {
-      $field_type_info = \Drupal::service('plugin.manager.field.type')->getDefinition(($context['field']['type']));
-
-      // If the field has specified a specific formatter to be used by default
-      // with tokens, use that, otherwise use the default formatter.
-      $formatter = !empty($field_type_info['default_token_formatter']) ? $field_type_info['default_token_formatter'] : $field_type_info['default_formatter'];
-
-      // Now that we have a formatter, fill in all the settings.
-      $display['type'] = $formatter;
-      $formatter_info = \Drupal::service('plugin.manager.field.formatter')->getDefinition($display['type']);
-      $display['settings'] = isset($formatter_info['settings']) ? $formatter_info['settings'] : array();
-      $display['settings']['label'] = 'hidden';
-      $display['module'] = $formatter_info['module'];
+      $fileTypeManager = \Drupal::service('plugin.manager.field.field_type');
+      $formatterManager = \Drupal::service('plugin.manager.field.formatter');
+
+      foreach ($entity_display->getComponents() as $field_name => $field_info) {
+        $field_type_info = FieldStorageConfig::loadByName($context['entity_type'], $field_name);
+        if (empty($field_type_info)) {
+          continue;
+        }
+        $field_type_definition = $fileTypeManager->getDefinition($field_type_info->getType());
+
+        // If the field has specified a specific formatter to be used by default
+        // with tokens, use that, otherwise use the default formatter.
+        $formatter = !empty($field_type_definition['default_token_formatter']) ? $field_type_definition['default_token_formatter'] : $field_type_definition['default_formatter'];
+
+        // Now that we have a formatter, fill in all the settings.
+        $formatter_info = $formatterManager->getDefinition($formatter);
+
+        $field_info['type'] = $formatter;
+        $field_info['settings'] = isset($formatter_info['settings']) ? $formatter_info['settings'] : array();
+        $field_info['settings']['label'] = 'hidden';
+
+        $entity_display->setComponent($field_name, $field_info);
+      }
     }
   }
 }