@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 181 lines 6.03 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import sys
5 from dataclasses import dataclass
6 from enum import Enum
7 from pathlib import Path
8
9
10 @dataclass(frozen=True)
11 class HeaderStyle:
12 """Comment syntax and preamble handling for an SPDX header."""
13
14 comment_prefix: str
15 preserve_shebang: bool = False
16
17
18 class SPDXStatus(Enum):
19 """SPDX header status enumeration."""
20
21 EMPTY = "empty" # empty __init__.py
22 COMPLETE = "complete"
23 MISSING_LICENSE = "missing_license" # Only has copyright line
24 MISSING_COPYRIGHT = "missing_copyright" # Only has license line
25 MISSING_BOTH = "missing_both" # Completely missing
26
27
28 LICENSE_TEXT = "SPDX-License-Identifier: Apache-2.0"
29 COPYRIGHT_TEXT = "SPDX-FileCopyrightText: Copyright contributors to the vLLM project"
30 FILE_STYLES = {
31 ".py": HeaderStyle("#", preserve_shebang=True),
32 ".rs": HeaderStyle("//"),
33 ".proto": HeaderStyle("//"),
34 }
35
36
37 def file_style(file_path):
38 """Return the declared header style for a file."""
39 suffix = Path(file_path).suffix
40 try:
41 return FILE_STYLES[suffix]
42 except KeyError:
43 raise ValueError(f"Unsupported file type: {file_path}") from None
44
45
46 def spdx_header(style):
47 """Return the SPDX header for a file style."""
48 license_line = f"{style.comment_prefix} {LICENSE_TEXT}"
49 copyright_line = f"{style.comment_prefix} {COPYRIGHT_TEXT}"
50 return license_line, copyright_line
51
52
53 def header_insertion_index(style, lines):
54 """Return the line index where a missing header should be inserted."""
55 if style.preserve_shebang and lines and lines[0].startswith("#!"):
56 return 1
57 return 0
58
59
60 def check_spdx_header_status(file_path):
61 """Check SPDX header status of the file."""
62 license_line, copyright_line = spdx_header(file_style(file_path))
63 with open(file_path, encoding="UTF-8") as file:
64 lines = file.readlines()
65 if not lines:
66 # Empty file
67 return SPDXStatus.EMPTY
68
69 has_license = False
70 has_copyright = False
71
72 # Check all lines for SPDX headers (not just the first two)
73 for raw_line in lines:
74 line = raw_line.strip()
75 if line == license_line:
76 has_license = True
77 elif line == copyright_line:
78 has_copyright = True
79
80 # Determine status based on what we found
81 if has_license and has_copyright:
82 return SPDXStatus.COMPLETE
83 elif has_license and not has_copyright:
84 # Only has license line
85 return SPDXStatus.MISSING_COPYRIGHT
86 elif not has_license and has_copyright:
87 # Only has copyright line
88 return SPDXStatus.MISSING_LICENSE
89 else:
90 # Completely missing both lines
91 return SPDXStatus.MISSING_BOTH
92
93
94 def add_header(file_path, status):
95 """Add or supplement SPDX header based on status."""
96 style = file_style(file_path)
97 license_line, copyright_line = spdx_header(style)
98 full_spdx_header = f"{license_line}\n{copyright_line}"
99 with open(file_path, "r+", encoding="UTF-8") as file:
100 lines = file.readlines()
101 file.seek(0, 0)
102 file.truncate()
103
104 if status == SPDXStatus.MISSING_BOTH:
105 # Completely missing, add complete header
106 insertion_index = header_insertion_index(style, lines)
107 file.writelines(lines[:insertion_index])
108 file.write(full_spdx_header + "\n")
109 remaining_lines = lines[insertion_index:]
110 if remaining_lines and remaining_lines[0].strip():
111 file.write("\n")
112 file.writelines(remaining_lines)
113
114 elif status == SPDXStatus.MISSING_COPYRIGHT:
115 # Only has license line, need to add copyright line
116 # Find the license line and add copyright line after it
117 for i, line in enumerate(lines):
118 if line.strip() == license_line:
119 # Insert copyright line after license line
120 lines.insert(
121 i + 1,
122 f"{copyright_line}\n",
123 )
124 break
125
126 file.writelines(lines)
127
128 elif status == SPDXStatus.MISSING_LICENSE:
129 # Only has copyright line, need to add license line
130 # Find the copyright line and add license line before it
131 for i, line in enumerate(lines):
132 if line.strip() == copyright_line:
133 # Insert license line before copyright line
134 lines.insert(i, f"{license_line}\n")
135 break
136 file.writelines(lines)
137
138
139 def main():
140 """Main function."""
141 files_missing_both = []
142 files_missing_copyright = []
143 files_missing_license = []
144
145 for file_path in sys.argv[1:]:
146 status = check_spdx_header_status(file_path)
147
148 if status == SPDXStatus.MISSING_BOTH:
149 files_missing_both.append(file_path)
150 elif status == SPDXStatus.MISSING_COPYRIGHT:
151 files_missing_copyright.append(file_path)
152 elif status == SPDXStatus.MISSING_LICENSE:
153 files_missing_license.append(file_path)
154 else:
155 continue
156
157 # Collect all files that need fixing
158 all_files_to_fix = (
159 files_missing_both + files_missing_copyright + files_missing_license
160 )
161 if all_files_to_fix:
162 print("The following files are missing the SPDX header:")
163 if files_missing_both:
164 for file_path in files_missing_both:
165 print(f" {file_path}")
166 add_header(file_path, SPDXStatus.MISSING_BOTH)
167
168 if files_missing_copyright:
169 for file_path in files_missing_copyright:
170 print(f" {file_path}")
171 add_header(file_path, SPDXStatus.MISSING_COPYRIGHT)
172 if files_missing_license:
173 for file_path in files_missing_license:
174 print(f" {file_path}")
175 add_header(file_path, SPDXStatus.MISSING_LICENSE)
176
177 sys.exit(1 if all_files_to_fix else 0)
178
179
180 if __name__ == "__main__":
181 main()