1 /*
2  *
3  *  * Copyright 2019-2020 the original author or authors.
4  *  *
5  *  * Licensed under the Apache License, Version 2.0 (the "License");
6  *  * you may not use this file except in compliance with the License.
7  *  * You may obtain a copy of the License at
8  *  *
9  *  *      https://www.apache.org/licenses/LICENSE-2.0
10  *  *
11  *  * Unless required by applicable law or agreed to in writing, software
12  *  * distributed under the License is distributed on an "AS IS" BASIS,
13  *  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14  *  * See the License for the specific language governing permissions and
15  *  * limitations under the License.
16  *
17  */
18
19 package org.springdoc.core.converters;
20
21 import java.util.HashMap;
22 import java.util.Iterator;
23 import java.util.Map;
24
25 import com.fasterxml.jackson.databind.JavaType;
26 import io.swagger.v3.core.converter.AnnotatedType;
27 import io.swagger.v3.core.converter.ModelConverter;
28 import io.swagger.v3.core.converter.ModelConverterContext;
29 import io.swagger.v3.core.util.Json;
30 import io.swagger.v3.oas.models.media.Schema;
31
32 public class AdditionalModelsConverter implements ModelConverter {
33
34     private static final Map<Class, Class> modelToClassMap = new HashMap();
35
36     private static final Map<Class, Schema> modelToSchemaMap = new HashMap();
37
38     public static void replaceWithClass(Class source, Class target) {
39         modelToClassMap.put(source, target);
40     }
41
42     public static void replaceWithSchema(Class source, Schema target) {
43         modelToSchemaMap.put(source, target);
44     }
45
46     public static Class getReplacement(Class clazz) {
47         return modelToClassMap.getOrDefault(clazz, clazz);
48     }
49
50     @Override
51     public Schema resolve(AnnotatedType type, ModelConverterContext context, Iterator<ModelConverter> chain) {
52         JavaType javaType = Json.mapper().constructType(type.getType());
53         if (javaType != null) {
54             Class<?> cls = javaType.getRawClass();
55             if (modelToSchemaMap.containsKey(cls))
56                 return modelToSchemaMap.get(cls);
57             if (modelToClassMap.containsKey(cls))
58                 type = new AnnotatedType(modelToClassMap.get(cls)).resolveAsRef(true);
59         }
60         return (chain.hasNext()) ? chain.next().resolve(type, context, chain) : null;
61     }
62
63 }