使用JavaPoet在代码生成中动态处理返回类型

核心思路 #

预定义Java内置类型（基本类型+常用内置引用类）的映射表，通过查表快速判断输入类型是否为内置类，未匹配到的则按自定义类处理，兼顾用户输入的简洁性和代码生成的准确性。

具体实现步骤 #

1. 构建内置类型映射表 #

创建一个Map存储用户输入的简单类型名与JavaPoet对应TypeName的映射，覆盖所有基本类型和常用内置引用类型：

import com.squareup.javapoet.ClassName;
import com.squareup.javapoet.TypeName;
import java.util.HashMap;
import java.util.Map;
import java.util.List;
import java.util.Map;
import java.util.Set;

private static final Map<String, TypeName> BUILT_IN_TYPE_MAP = new HashMap<>();

static {
    // 注册基本类型
    BUILT_IN_TYPE_MAP.put("int", TypeName.INT);
    BUILT_IN_TYPE_MAP.put("long", TypeName.LONG);
    BUILT_IN_TYPE_MAP.put("float", TypeName.FLOAT);
    BUILT_IN_TYPE_MAP.put("double", TypeName.DOUBLE);
    BUILT_IN_TYPE_MAP.put("boolean", TypeName.BOOLEAN);
    BUILT_IN_TYPE_MAP.put("byte", TypeName.BYTE);
    BUILT_IN_TYPE_MAP.put("char", TypeName.CHAR);
    BUILT_IN_TYPE_MAP.put("short", TypeName.SHORT);

    // 注册常用内置引用类型
    BUILT_IN_TYPE_MAP.put("String", ClassName.get(String.class));
    BUILT_IN_TYPE_MAP.put("Integer", ClassName.get(Integer.class));
    BUILT_IN_TYPE_MAP.put("List", ClassName.get(List.class));
    BUILT_IN_TYPE_MAP.put("Map", ClassName.get(Map.class));
    BUILT_IN_TYPE_MAP.put("Set", ClassName.get(Set.class));
    // 可根据需求添加更多内置类，如ArrayList、HashMap等
}

2. 实现类型解析工具方法 #

写一个工具方法，接收用户输入的returnType字符串，返回对应的JavaPoetTypeName：

import java.util.Arrays;

public static TypeName resolveReturnType(String returnType) {
    // 处理输入首尾空格，可选添加大小写兼容（比如将输入转成小写匹配）
    String normalizedType = returnType.trim();
    
    // 优先匹配内置类型
    TypeName builtInType = BUILT_IN_TYPE_MAP.get(normalizedType);
    if (builtInType != null) {
        return builtInType;
    }

    // 处理自定义类：支持全限定类名（如com.example.User）和简单类名
    if (normalizedType.contains(".")) {
        // 拆分包名和类名
        String[] nameParts = normalizedType.split("\\.");
        String packageName = String.join(".", Arrays.copyOfRange(nameParts, 0, nameParts.length - 1));
        String className = nameParts[nameParts.length - 1];
        return ClassName.get(packageName, className);
    }

    // 简单类名的自定义类，默认使用当前生成类的包名（可根据上下文调整）
    // 示例中用默认包，实际开发中可传入生成类的包名作为参数
    return ClassName.get("", normalizedType);
}

3. 扩展处理特殊类型（可选） #

如果需要支持数组、泛型等复杂类型，可以在工具方法中添加额外逻辑：

• 数组类型：比如用户输入int[]，可先提取基础类型int，再调用TypeName.arrayOf(resolveReturnType("int"))生成数组类型；
• 泛型类型：比如用户输入List<String>，需要解析泛型参数，使用ParameterizedTypeName.get((ClassName) resolveReturnType("List"), resolveReturnType("String"))生成泛型类型。

使用示例 #

在生成方法时直接调用工具方法：

// 用户输入的参数
String methodName = "getUserInfo";
String returnTypeInput = "String";
// 解析返回类型
TypeName returnType = resolveReturnType(returnTypeInput);

// 生成方法
MethodSpec methodSpec = MethodSpec.methodBuilder(methodName)
        .addModifiers(Modifier.PUBLIC)
        .returns(returnType)
        // 添加方法体等逻辑
        .build();

内容的提问来源于stack exchange，提问作者Hiroshking