TransformerEngine

1
#!/usr/bin/env python
2
# coding: utf-8
3

4
# Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
5
#
6
# See LICENSE for license information.
7

8

9
import os
10
import sys
11
import json
12
import datetime
13

14
if len(sys.argv) < 2:
15
    print("Usage: python copyright_checker.py <path>")
16

17
path = sys.argv[1]
18

19
config_path = os.path.dirname(os.path.realpath(__file__)) + "/config.json"
20

21
class bcolors:
22
    OKGREEN = '\033[92m'
23
    WARNING = '\033[93m'
24
    FAIL = '\033[91m'
25
    ENDC = '\033[0m'
26

27
def print_ok(msg):
28
    print(f"{bcolors.OKGREEN}{msg}{bcolors.ENDC}")
29

30
def print_fail(msg):
31
    print(f"{bcolors.FAIL}{msg}{bcolors.ENDC}")
32

33
def print_warn(msg):
34
    print(f"{bcolors.WARNING}{msg}{bcolors.ENDC}")
35

36
with open(config_path, "r") as f:
37
    c = json.load(f)
38
    current_year = datetime.date.today().year
39
    if c["initial_year"] == current_year:
40
        year_string = str(current_year)
41
    else:
42
        year_string = str(c["initial_year"]) + "-" + str(current_year)
43
    copyright_string = c["copyright"].replace("<YEAR>", year_string)
44
    license = c["license"].split('\n')
45
    excludes = c["exclude"]
46
    root_path = os.path.abspath(path)
47
    copyright_only = c["copyright_only"]
48
    exclude_copyright = c["exclude_copyright"]
49

50
has_gitignore = os.path.exists(root_path + "/.gitignore")
51

52
def strip_star_slash(s):
53
    ret = s
54
    if ret.startswith('*'):
55
        ret = ret[1:]
56
    if ret.endswith('/'):
57
        ret = ret[:-1]
58
    return ret
59

60
if has_gitignore:
61
    with open(root_path + "/.gitignore", "r") as f:
62
        for line in f.readlines():
63
            excludes.append(strip_star_slash(line.strip()))
64

65
def get_file_type(path):
66
    ext = {"c": ["c", "cpp", "cu", "h", "cuh"],
67
           "py": ["py"],
68
           "rst": ["rst"],
69
           "txt": ["txt"],
70
           "cfg": ["cfg"],
71
           "sh":  ["sh"],
72
           "md":  ["md"],
73
          }
74
    tmp = path.split(".")
75
    for filetype, ext_list in ext.items():
76
        if tmp[-1] in ext_list:
77
            return filetype
78
    return "unknown"
79

80
success = True
81

82
def check_file(path):
83
    global success
84
    N = 10
85
    ftype = get_file_type(path)
86
    if ftype == "unknown":
87
        print_warn("Unknown filetype")
88
        return
89
    check_copyright = True
90
    for e in exclude_copyright:
91
        if path.endswith(e):
92
            check_copyright = False
93
    with open(path, "r") as f:
94
        copyright_found = False
95
        license_found = True
96
        try:
97
            if check_copyright:
98
                for _ in range(N):
99
                    line = f.readline()
100
                    if line.find(copyright_string) != -1:
101
                        copyright_found = True
102
                        break
103
            if not copyright_only:
104
                first_license_line = True
105
                for l in license:
106
                    if first_license_line:
107
                        # may skip some lines
108
                        first_license_line = False
109
                        for _ in range(N):
110
                            line = f.readline()
111
                            if line.find(l) != -1:
112
                                break
113
                    else:
114
                        line = f.readline()
115
                    if line.find(l) == -1:
116
                        license_found = False
117
                        break
118
        except:
119
            pass
120
        finally:
121
            if not copyright_found:
122
                print_fail("No copyright found!")
123
                success = False
124
            if not license_found:
125
                print_fail("No license found!")
126
                success = False
127
            if copyright_found and license_found:
128
                print_ok("OK")
129

130
for root, dirs, files in os.walk(root_path):
131
    print(f"Entering {root}")
132
    hidden = [d for d in dirs if d.startswith('.')] + [f for f in files if f.startswith('.')]
133
    all_excludes = excludes + hidden
134
    to_remove = []
135
    for d in dirs:
136
        d_path = root + "/" + d
137
        for e in all_excludes:
138
            if d_path.endswith(e):
139
                to_remove.append(d)
140
    for f in files:
141
        f_path = root + "/" + f
142
        for e in all_excludes:
143
            if f_path.endswith(e):
144
                to_remove.append(f)
145
    for d in to_remove:
146
        if d in dirs:
147
            dirs.remove(d)
148
        if d in files:
149
            files.remove(d)
150
    for filename in files:
151
        print(f"Checking {filename}")
152
        check_file(os.path.abspath(root + "/" + filename))
153

154
if not success:
155
    raise Exception("Some copyrights/licenses are missing!")
156