From e9cca364b7a4b0304510ebc9155e6bb44dbae5b5 Mon Sep 17 00:00:00 2001
From: Hunter Mellema <124718352+hpmellema@users.noreply.github.com>
Date: Wed, 28 Feb 2024 12:13:23 -0700
Subject: [PATCH] Add data-shape-only visitor class (#2168)

---
 .../smithy/model/shapes/ShapeVisitor.java     | 26 +++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/smithy-model/src/main/java/software/amazon/smithy/model/shapes/ShapeVisitor.java b/smithy-model/src/main/java/software/amazon/smithy/model/shapes/ShapeVisitor.java
index 6a62ba74709..b36c616f20d 100644
--- a/smithy-model/src/main/java/software/amazon/smithy/model/shapes/ShapeVisitor.java
+++ b/smithy-model/src/main/java/software/amazon/smithy/model/shapes/ShapeVisitor.java
@@ -199,4 +199,30 @@ public R timestampShape(TimestampShape shape) {
             return getDefault(shape);
         }
     }
+
+    /**
+     * Creates {@link ShapeVisitor} that only requires implementation of
+     * all data shape branches, but does not support service shapes.
+     *
+     * @param <R> Return type.
+     */
+    abstract class DataShapeVisitor<R> implements ShapeVisitor<R> {
+        @Override
+        public R operationShape(OperationShape shape) {
+            throw new IllegalArgumentException("DataShapeVisitor cannot be use to visit "
+                    + "Operation Shapes. Attempted to visit: " + shape);
+        }
+
+        @Override
+        public R resourceShape(ResourceShape shape) {
+            throw new IllegalArgumentException("DataShapeVisitor cannot be use to visit "
+                    + "Resource Shapes. Attempted to visit: " + shape);
+        }
+
+        @Override
+        public R serviceShape(ServiceShape shape) {
+            throw new IllegalArgumentException("DataShapeVisitor cannot be use to visit "
+                + "Service Shapes. Attempted to visit: " + shape);
+        }
+    }
 }