com.squareup.wire.sample.CodegenSample.java Source code

Introduction

Here is the source code for com.squareup.wire.sample.CodegenSample.java

Source

/*
 * Copyright (C) 2016 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.squareup.wire.sample;

import com.google.common.base.Stopwatch;
import com.google.common.collect.ImmutableSet;
import com.squareup.javapoet.ClassName;
import com.squareup.javapoet.JavaFile;
import com.squareup.javapoet.TypeSpec;
import com.squareup.wire.java.JavaGenerator;
import com.squareup.wire.schema.EnumType;
import com.squareup.wire.schema.IdentifierSet;
import com.squareup.wire.schema.Location;
import com.squareup.wire.schema.MessageType;
import com.squareup.wire.schema.ProtoFile;
import com.squareup.wire.schema.Schema;
import com.squareup.wire.schema.SchemaLoader;
import com.squareup.wire.schema.Service;
import com.squareup.wire.schema.Type;
import java.io.File;
import java.io.IOException;

/**
 * Reads in .proto files, links em together, prunes em, and emits both types and services. This
 * can be used to either complement or replace {@code WireCompiler}.
 */
public final class CodegenSample {
    private final Log log;
    private final ImmutableSet<String> sources;
    private final ImmutableSet<String> protos;
    private final String generatedSourceDirectory;
    private final IdentifierSet identifierSet;

    public CodegenSample(Log log, ImmutableSet<String> sources, ImmutableSet<String> protos,
            String generatedSourceDirectory, IdentifierSet identifierSet) {
        this.log = log;
        this.sources = sources;
        this.protos = protos;
        this.generatedSourceDirectory = generatedSourceDirectory;
        this.identifierSet = identifierSet;
    }

    public void execute() throws IOException {
        Schema schema = loadSchema();

        if (!identifierSet.isEmpty()) {
            schema = retainRoots(schema);
        }

        JavaGenerator javaGenerator = JavaGenerator.get(schema);
        ServiceGenerator serviceGenerator = new ServiceGenerator(javaGenerator);

        for (ProtoFile protoFile : schema.protoFiles()) {
            for (Type type : protoFile.types()) {
                Stopwatch stopwatch = Stopwatch.createStarted();
                ClassName javaTypeName = (ClassName) javaGenerator.typeName(type.type());
                TypeSpec typeSpec = type instanceof MessageType ? javaGenerator.generateMessage((MessageType) type)
                        : javaGenerator.generateEnum((EnumType) type);
                writeJavaFile(javaTypeName, typeSpec, type.location(), stopwatch);
            }

            for (Service service : protoFile.services()) {
                Stopwatch stopwatch = Stopwatch.createStarted();
                ClassName javaTypeName = (ClassName) javaGenerator.typeName(service.type());
                TypeSpec typeSpec = serviceGenerator.api(service);
                writeJavaFile(javaTypeName, typeSpec, service.location(), stopwatch);
            }
        }
    }

    private Schema loadSchema() throws IOException {
        Stopwatch stopwatch = Stopwatch.createStarted();
        SchemaLoader loader = new SchemaLoader();
        for (String source : sources) {
            loader.addSource(new File(source));
        }
        for (String proto : protos) {
            loader.addProto(proto);
        }
        Schema schema = loader.load();
        log.info("Loaded %s proto files in %s", schema.protoFiles().size(), stopwatch);
        return schema;
    }

    private Schema retainRoots(Schema schema) {
        Stopwatch stopwatch = Stopwatch.createStarted();
        int oldSize = countTypes(schema);

        Schema prunedSchema = schema.prune(identifierSet);
        int newSize = countTypes(prunedSchema);

        log.info("Pruned schema from %s types to %s types in %s", oldSize, newSize, stopwatch);

        return prunedSchema;
    }

    private void writeJavaFile(ClassName javaTypeName, TypeSpec typeSpec, Location location, Stopwatch stopwatch)
            throws IOException {
        JavaFile.Builder builder = JavaFile.builder(javaTypeName.packageName(), typeSpec)
                .addFileComment("Code generated by $L, do not edit.", CodegenSample.class.getName());
        if (location != null) {
            builder.addFileComment("\nSource file: $L", location.path());
        }
        JavaFile javaFile = builder.build();
        try {
            javaFile.writeTo(new File(generatedSourceDirectory));
        } catch (IOException e) {
            throw new IOException("Failed to write " + javaFile.packageName + "." + javaFile.typeSpec.name + " to "
                    + generatedSourceDirectory, e);
        }
        log.info("Generated %s in %s", javaTypeName, stopwatch);
    }

    private int countTypes(Schema prunedSchema) {
        int result = 0;
        for (ProtoFile protoFile : prunedSchema.protoFiles()) {
            result += protoFile.types().size();
        }
        return result;
    }

    interface Log {
        void info(String format, Object... args);
    }
}