com.itfsw.mybatis.generator.plugins.utils.enhanced.TemplateCommentGenerator.java Source code

Java tutorial

Introduction

Here is the source code for com.itfsw.mybatis.generator.plugins.utils.enhanced.TemplateCommentGenerator.java

Source

/*
 * Copyright (c) 2017.
 *
 * 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.itfsw.mybatis.generator.plugins.utils.enhanced;

import freemarker.template.Configuration;
import freemarker.template.Template;
import org.dom4j.Document;
import org.dom4j.Element;
import org.dom4j.io.SAXReader;
import org.mybatis.generator.api.CommentGenerator;
import org.mybatis.generator.api.IntrospectedColumn;
import org.mybatis.generator.api.IntrospectedTable;
import org.mybatis.generator.api.MyBatisGenerator;
import org.mybatis.generator.api.dom.java.*;
import org.mybatis.generator.api.dom.xml.TextElement;
import org.mybatis.generator.api.dom.xml.XmlElement;
import org.mybatis.generator.config.MergeConstants;
import org.mybatis.generator.config.PropertyRegistry;
import org.mybatis.generator.internal.util.StringUtility;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.xml.bind.DatatypeConverter;
import java.io.File;
import java.io.InputStream;
import java.io.StringWriter;
import java.util.*;

import static org.mybatis.generator.internal.util.StringUtility.isTrue;

/**
 * ---------------------------------------------------------------------------
 * ??
 * ---------------------------------------------------------------------------
 * @author: hewei
 * @time:2017/6/8 13:21
 * ---------------------------------------------------------------------------
 */
public class TemplateCommentGenerator implements CommentGenerator {
    protected static final Logger logger = LoggerFactory.getLogger(TemplateCommentGenerator.class); // 

    private Map<EnumNode, Template> templates = new HashMap<>(); // ?

    private boolean suppressDate = false;

    private boolean suppressAllComments = false;

    /**
     * 
     * @param templatePath ?
     * @param useForDefault Comment??
     */
    public TemplateCommentGenerator(String templatePath, boolean useForDefault) {
        try {
            Document doc = null;
            if (useForDefault) {
                InputStream inputStream = this.getClass().getClassLoader().getResourceAsStream(templatePath);
                doc = new SAXReader().read(inputStream);
                inputStream.close();
            } else {
                File file = new File(templatePath);
                if (file.exists()) {
                    doc = new SAXReader().read(file);
                } else {
                    logger.error("?:" + templatePath);
                }
            }

            // ??comment 
            if (doc != null) {
                for (EnumNode node : EnumNode.values()) {
                    Element element = doc.getRootElement().elementByID(node.value());
                    if (element != null) {
                        Configuration cfg = new Configuration(Configuration.VERSION_2_3_26);
                        // ?
                        Template template = new Template(node.value(), element.getText(), cfg);
                        templates.put(node, template);
                    }
                }
            }
        } catch (Exception e) {
            logger.error("?XML??", e);
        }
    }

    /**
     * ?
     * @param map ??
     * @param node  ID
     * @return
     */
    private String[] getComments(Map<String, Object> map, EnumNode node) {
        // 1. ??
        try {
            StringWriter stringWriter = new StringWriter();
            Template template = templates.get(node);
            if (template != null) {
                template.process(map, stringWriter);

                String comment = stringWriter.toString();
                stringWriter.close();

                // ??
                return comment.replaceFirst("^[\\s\\t\\r\\n]*", "").replaceFirst("[\\s\\t\\r\\n]*$", "")
                        .split("\n");
            }

        } catch (Exception e) {
            logger.error("freemarker ??", e);
        }

        return null;
    }

    /**
     * 
     *
     * @param javaElement
     * @param map
     * @param node
     */
    private void addJavaElementComment(JavaElement javaElement, Map<String, Object> map, EnumNode node) {
        // ?
        String[] comments = getComments(map, node);
        if (comments != null) {
            // 
            if (comments.length == 1 && !StringUtility.stringHasValue(comments[0])) {
                return;
            }
            // 
            for (String comment : comments) {
                javaElement.addJavaDocLine(comment);
            }
        }
    }

    /**
     * 
     *
     * @param compilationUnit
     * @param map
     * @param node
     */
    private void addCompilationUnitComment(CompilationUnit compilationUnit, Map<String, Object> map,
            EnumNode node) {
        // ?
        String[] comments = getComments(map, node);
        if (comments != null) {
            // 
            if (comments.length == 1 && !StringUtility.stringHasValue(comments[0])) {
                return;
            }
            // 
            for (String comment : comments) {
                compilationUnit.addFileCommentLine(comment);
            }
        }
    }

    /**
     * 
     *
     * @param xmlElement
     * @param map
     * @param node
     */
    private void addXmlElementComment(XmlElement xmlElement, Map<String, Object> map, EnumNode node) {
        // ?
        String[] comments = getComments(map, node);
        if (comments != null) {
            // 
            if (comments.length == 1 && !StringUtility.stringHasValue(comments[0])) {
                return;
            }
            // 
            for (String comment : comments) {
                xmlElement.addElement(new TextElement(comment));
            }
        }
    }

    /**
     * Adds properties for this instance from any properties configured in the
     * CommentGenerator configuration.
     *
     * This method will be called before any of the other methods.
     *
     * @param properties
     *            All properties from the configuration
     */
    @Override
    public void addConfigurationProperties(Properties properties) {
        suppressDate = isTrue(properties.getProperty(PropertyRegistry.COMMENT_GENERATOR_SUPPRESS_DATE));

        suppressAllComments = isTrue(
                properties.getProperty(PropertyRegistry.COMMENT_GENERATOR_SUPPRESS_ALL_COMMENTS));
    }

    /**
     * This method should add a Javadoc comment to the specified field. The field is related to the specified table and
     * is used to hold the value of the specified column.
     * <p>
     *
     * <b>Important:</b> This method should add a the nonstandard JavaDoc tag "@mbg.generated" to the comment. Without
     * this tag, the Eclipse based Java merge feature will fail.
     *
     * @param field
     *            the field
     * @param introspectedTable
     *            the introspected table
     * @param introspectedColumn
     *            the introspected column
     */
    @Override
    public void addFieldComment(Field field, IntrospectedTable introspectedTable,
            IntrospectedColumn introspectedColumn) {
        Map<String, Object> map = new HashMap<>();
        map.put("mgb", MergeConstants.NEW_ELEMENT_TAG);
        map.put("field", field);
        map.put("introspectedTable", introspectedTable);
        map.put("introspectedColumn", introspectedColumn);

        // 
        addJavaElementComment(field, map, EnumNode.ADD_FIELD_COMMENT);
    }

    /**
     * Adds the field comment.
     *
     * @param field
     *            the field
     * @param introspectedTable
     *            the introspected table
     */
    @Override
    public void addFieldComment(Field field, IntrospectedTable introspectedTable) {
        Map<String, Object> map = new HashMap<>();
        map.put("mgb", MergeConstants.NEW_ELEMENT_TAG);
        map.put("field", field);
        map.put("introspectedTable", introspectedTable);

        // 
        addJavaElementComment(field, map, EnumNode.ADD_FIELD_COMMENT);
    }

    /**
     * Adds a comment for a model class.  The Java code merger should
     * be notified not to delete the entire class in case any manual
     * changes have been made.  So this method will always use the
     * "do not delete" annotation.
     *
     * Because of difficulties with the Java file merger, the default implementation
     * of this method should NOT add comments.  Comments should only be added if
     * specifically requested by the user (for example, by enabling table remark comments).
     *
     * @param topLevelClass
     *            the top level class
     * @param introspectedTable
     *            the introspected table
     */
    @Override
    public void addModelClassComment(TopLevelClass topLevelClass, IntrospectedTable introspectedTable) {
        Map<String, Object> map = new HashMap<>();
        map.put("mgb", MergeConstants.NEW_ELEMENT_TAG);
        map.put("topLevelClass", topLevelClass);
        map.put("introspectedTable", introspectedTable);

        // 
        addJavaElementComment(topLevelClass, map, EnumNode.ADD_MODEL_CLASS_COMMENT);
    }

    /**
     * Adds the inner class comment.
     *
     * @param innerClass
     *            the inner class
     * @param introspectedTable
     *            the introspected table
     */
    @Override
    public void addClassComment(InnerClass innerClass, IntrospectedTable introspectedTable) {
        if (innerClass instanceof InnerInterfaceWrapperToInnerClass) {
            InnerInterface innerInterface = ((InnerInterfaceWrapperToInnerClass) innerClass).getInnerInterface();

            Map<String, Object> map = new HashMap<>();
            map.put("mgb", MergeConstants.NEW_ELEMENT_TAG);
            map.put("innerInterface", innerInterface);
            map.put("introspectedTable", introspectedTable);

            // 
            addJavaElementComment(innerInterface, map, EnumNode.ADD_INTERFACE_COMMENT);
        } else {
            Map<String, Object> map = new HashMap<>();
            map.put("mgb", MergeConstants.NEW_ELEMENT_TAG);
            map.put("innerClass", innerClass);
            map.put("introspectedTable", introspectedTable);

            // 
            addJavaElementComment(innerClass, map, EnumNode.ADD_CLASS_COMMENT);
        }
    }

    /**
     * Adds the inner class comment.
     *
     * @param innerClass
     *            the inner class
     * @param introspectedTable
     *            the introspected table
     * @param markAsDoNotDelete
     *            the mark as do not delete
     */
    @Override
    public void addClassComment(InnerClass innerClass, IntrospectedTable introspectedTable,
            boolean markAsDoNotDelete) {
        Map<String, Object> map = new HashMap<>();
        map.put("mgb", MergeConstants.NEW_ELEMENT_TAG);
        map.put("innerClass", innerClass);
        map.put("introspectedTable", introspectedTable);
        map.put("markAsDoNotDelete", markAsDoNotDelete);

        // 
        addJavaElementComment(innerClass, map, EnumNode.ADD_CLASS_COMMENT);
    }

    /**
     * Adds the enum comment.
     *
     * @param innerEnum
     *            the inner enum
     * @param introspectedTable
     *            the introspected table
     */
    @Override
    public void addEnumComment(InnerEnum innerEnum, IntrospectedTable introspectedTable) {
        Map<String, Object> map = new HashMap<>();
        map.put("mgb", MergeConstants.NEW_ELEMENT_TAG);
        map.put("innerEnum", innerEnum);
        map.put("introspectedTable", introspectedTable);

        // 
        addJavaElementComment(innerEnum, map, EnumNode.ADD_ENUM_COMMENT);
    }

    /**
     * Adds the getter comment.
     *
     * @param method
     *            the method
     * @param introspectedTable
     *            the introspected table
     * @param introspectedColumn
     *            the introspected column
     */
    @Override
    public void addGetterComment(Method method, IntrospectedTable introspectedTable,
            IntrospectedColumn introspectedColumn) {
        Map<String, Object> map = new HashMap<>();
        map.put("mgb", MergeConstants.NEW_ELEMENT_TAG);
        map.put("method", method);
        map.put("introspectedTable", introspectedTable);
        map.put("introspectedColumn", introspectedColumn);

        // 
        addJavaElementComment(method, map, EnumNode.ADD_GETTER_COMMENT);
    }

    /**
     * Adds the setter comment.
     *
     * @param method
     *            the method
     * @param introspectedTable
     *            the introspected table
     * @param introspectedColumn
     *            the introspected column
     */
    @Override
    public void addSetterComment(Method method, IntrospectedTable introspectedTable,
            IntrospectedColumn introspectedColumn) {
        Map<String, Object> map = new HashMap<>();
        map.put("mgb", MergeConstants.NEW_ELEMENT_TAG);
        map.put("method", method);
        map.put("introspectedTable", introspectedTable);
        map.put("introspectedColumn", introspectedColumn);

        // 
        addJavaElementComment(method, map, EnumNode.ADD_SETTER_COMMENT);
    }

    /**
     * Adds the general method comment.
     *
     * @param method
     *            the method
     * @param introspectedTable
     *            the introspected table
     */
    @Override
    public void addGeneralMethodComment(Method method, IntrospectedTable introspectedTable) {
        Map<String, Object> map = new HashMap<>();
        map.put("mgb", MergeConstants.NEW_ELEMENT_TAG);
        map.put("method", method);
        map.put("introspectedTable", introspectedTable);

        // 
        addJavaElementComment(method, map, EnumNode.ADD_GENERAL_METHOD_COMMENT);
    }

    /**
     * This method is called to add a file level comment to a generated java file. This method could be used to add a
     * general file comment (such as a copyright notice). However, note that the Java file merge function in Eclipse
     * does not deal with this comment. If you run the generator repeatedly, you will only retain the comment from the
     * initial run.
     * <p>
     *
     * The default implementation does nothing.
     *
     * @param compilationUnit
     *            the compilation unit
     */
    @Override
    public void addJavaFileComment(CompilationUnit compilationUnit) {
        Map<String, Object> map = new HashMap<>();
        map.put("mgb", MergeConstants.NEW_ELEMENT_TAG);
        map.put("compilationUnit", compilationUnit);

        // 
        addCompilationUnitComment(compilationUnit, map, EnumNode.ADD_JAVA_FILE_COMMENT);
    }

    /**
     * This method should add a suitable comment as a child element of the specified xmlElement to warn users that the
     * element was generated and is subject to regeneration.
     *
     * @param xmlElement
     *            the xml element
     */
    @Override
    public void addComment(XmlElement xmlElement) {
        Map<String, Object> map = new HashMap<>();
        map.put("mgb", MergeConstants.NEW_ELEMENT_TAG);
        map.put("xmlElement", xmlElement);

        // 
        addXmlElementComment(xmlElement, map, EnumNode.ADD_COMMENT);
    }

    /**
     * This method is called to add a comment as the first child of the root element. This method could be used to add a
     * general file comment (such as a copyright notice). However, note that the XML file merge function does not deal
     * with this comment. If you run the generator repeatedly, you will only retain the comment from the initial run.
     * <p>
     *
     * The default implementation does nothing.
     *
     * @param rootElement
     *            the root element
     */
    @Override
    public void addRootComment(XmlElement rootElement) {
        Map<String, Object> map = new HashMap<>();
        map.put("mgb", MergeConstants.NEW_ELEMENT_TAG);
        map.put("rootElement", rootElement);

        // 
        addXmlElementComment(rootElement, map, EnumNode.ADD_ROOT_COMMENT);
    }

    @Override
    public void addGeneralMethodAnnotation(Method method, IntrospectedTable introspectedTable,
            Set<FullyQualifiedJavaType> imports) {
        imports.add(new FullyQualifiedJavaType("javax.annotation.Generated")); //$NON-NLS-1$
        String comment = "Source Table: " + introspectedTable.getFullyQualifiedTable().toString(); //$NON-NLS-1$
        method.addAnnotation(getGeneratedAnnotation(comment));
    }

    @Override
    public void addGeneralMethodAnnotation(Method method, IntrospectedTable introspectedTable,
            IntrospectedColumn introspectedColumn, Set<FullyQualifiedJavaType> imports) {
        imports.add(new FullyQualifiedJavaType("javax.annotation.Generated")); //$NON-NLS-1$
        String comment = "Source field: " //$NON-NLS-1$
                + introspectedTable.getFullyQualifiedTable().toString() + "." //$NON-NLS-1$
                + introspectedColumn.getActualColumnName();
        method.addAnnotation(getGeneratedAnnotation(comment));
    }

    @Override
    public void addFieldAnnotation(Field field, IntrospectedTable introspectedTable,
            Set<FullyQualifiedJavaType> imports) {
        imports.add(new FullyQualifiedJavaType("javax.annotation.Generated")); //$NON-NLS-1$
        String comment = "Source Table: " + introspectedTable.getFullyQualifiedTable().toString(); //$NON-NLS-1$
        field.addAnnotation(getGeneratedAnnotation(comment));
    }

    @Override
    public void addFieldAnnotation(Field field, IntrospectedTable introspectedTable,
            IntrospectedColumn introspectedColumn, Set<FullyQualifiedJavaType> imports) {
        imports.add(new FullyQualifiedJavaType("javax.annotation.Generated")); //$NON-NLS-1$
        String comment = "Source field: " //$NON-NLS-1$
                + introspectedTable.getFullyQualifiedTable().toString() + "." //$NON-NLS-1$
                + introspectedColumn.getActualColumnName();
        field.addAnnotation(getGeneratedAnnotation(comment));
    }

    @Override
    public void addClassAnnotation(InnerClass innerClass, IntrospectedTable introspectedTable,
            Set<FullyQualifiedJavaType> imports) {
        imports.add(new FullyQualifiedJavaType("javax.annotation.Generated")); //$NON-NLS-1$
        String comment = "Source Table: " + introspectedTable.getFullyQualifiedTable().toString(); //$NON-NLS-1$
        innerClass.addAnnotation(getGeneratedAnnotation(comment));
    }

    private String getGeneratedAnnotation(String comment) {
        StringBuilder buffer = new StringBuilder();
        buffer.append("@Generated("); //$NON-NLS-1$
        if (suppressAllComments) {
            buffer.append('\"');
        } else {
            buffer.append("value=\""); //$NON-NLS-1$
        }

        buffer.append(MyBatisGenerator.class.getName());
        buffer.append('\"');

        if (!suppressDate && !suppressAllComments) {
            buffer.append(", date=\""); //$NON-NLS-1$
            buffer.append(DatatypeConverter.printDateTime(Calendar.getInstance()));
            buffer.append('\"');
        }

        if (!suppressAllComments) {
            buffer.append(", comments=\""); //$NON-NLS-1$
            buffer.append(comment);
            buffer.append('\"');
        }

        buffer.append(')');
        return buffer.toString();
    }

    /**
     * ?ID
     */
    public static enum EnumNode {
        ADD_COMMENT("addComment"), // Xml 
        ADD_ROOT_COMMENT("addRootComment"), // xml root 
        ADD_JAVA_FILE_COMMENT("addJavaFileComment"), // java 
        ADD_GENERAL_METHOD_COMMENT("addGeneralMethodComment"), // java 
        ADD_SETTER_COMMENT("addSetterComment"), // setter 
        ADD_GETTER_COMMENT("addGetterComment"), // getter ?
        ADD_ENUM_COMMENT("addEnumComment"), //  
        ADD_CLASS_COMMENT("addClassComment"), //  
        ADD_INTERFACE_COMMENT("addInterfaceComment"), // ? 
        ADD_MODEL_CLASS_COMMENT("addModelClassComment"), // model 
        ADD_FIELD_COMMENT("addFieldComment"), //  
        ;

        private final String value; // 

        /**
         * 
         * @param value
         */
        EnumNode(String value) {
            this.value = value;
        }

        /**
         * 
         *
         * @return
         */
        public String value() {
            return value;
        }
    }
}