From b2f8b96a0221e0e61f704c591b8fcc6fc1665cc6 Mon Sep 17 00:00:00 2001
From: eqy <eddieyan101@gmail.com>
Date: Thu, 23 May 2019 17:52:44 -0700
Subject: [PATCH] [LINT] handle more file types in ASF header (#3235)

* Update add_asf_header.py

* Update add_asf_header.py
---
 tests/lint/add_asf_header.py | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/tests/lint/add_asf_header.py b/tests/lint/add_asf_header.py
index 7e0352f4bc2c..1afb3a57b2f1 100644
--- a/tests/lint/add_asf_header.py
+++ b/tests/lint/add_asf_header.py
@@ -117,6 +117,7 @@
 """.strip()
 
 FMT_MAP = {
+    "sh" : header_pystyle,
     "cc" : header_cstyle,
     "h" : header_cstyle,
     "py" : header_pystyle,
@@ -128,6 +129,7 @@
     "cmake" : header_pystyle,
     "rst" : header_rststyle,
     "gradle" : header_groovystyle,
+    "xml": header_mdstyle,
 }
 
 def add_header(fname, header):
@@ -142,8 +144,23 @@ def add_header(fname, header):
         return
 
     with open(fname, "w") as outfile:
-        outfile.write(header + "\n\n")
-        outfile.write(orig)
+        skipline = False
+        lines = orig.split('\n')
+        ext = os.path.splitext(fname)[1][1:]
+        if ext == 'sh' and lines[0][:2] == '#!':
+            skipline = True
+        elif ext == 'xml' and lines[0][:2] == '<?':
+            skipline = True
+
+        if skipline:
+            outfile.write(lines[0] + "\n")
+            outfile.write(header + "\n\n")
+            outfile.write("\n".join(lines[1:]))
+            outfile.write(header + "\n\n")
+            outfile.write(orig)
+        else:
+            outfile.write(header + "\n\n")
+            outfile.write(orig)
     print("Add header to %s" % fname)
 
 
@@ -160,6 +177,8 @@ def main(args):
         suffix = fname.split(".")[-1]
         if suffix in FMT_MAP:
             add_header(fname, FMT_MAP[suffix])
+        elif os.path.basename(fname) == 'gradle.properties':
+            add_header(fname, FMT_MAP['h'])
         else:
             print("Cannot handle %s ..." % fname)