Skip to content

Commit 2c265b2

Browse files
author
Markus Wegmann
committed
Initial commit
0 parents  commit 2c265b2

File tree

6 files changed

+245
-0
lines changed

6 files changed

+245
-0
lines changed

‎.gitignore

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
.vscode/
2+
3+
*.egg-info
4+
*.tar.gz
5+
*.whl
6+
7+
build/

‎LICENSE

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
Copyright 2022 Markus Wegmann
2+
3+
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
4+
5+
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
6+
7+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

‎README.md

+60
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
# fix-protobuf-imports
2+
3+
This script will fix relative imports (from and to nested sub-directories) within compiled `*pb2.py` and `*pb2.pyi` Protobuf files/modules generated from `protoc --python_out --mypy_out`:
4+
5+
```bash
6+
fix-protobuf-imports /path/to/python_out/dir
7+
```
8+
9+
## When do I need to fix my imports?
10+
11+
E.g. you might have the following file/module structure:
12+
13+
- `./`
14+
- `a_pb2.py`
15+
- `b_pb2.py`
16+
- `./sub/`
17+
- `c_pb2.py`
18+
- `./nested/`
19+
- `d_pb2.py`
20+
- `__init__.py`
21+
- `__init__.py`
22+
- `__init__.py`
23+
24+
Now assume, `a_sub.proto` is importing `a.proto`, `b.proto` and `c.proto`.
25+
26+
`protoc` will generate the following import statements:
27+
28+
```python
29+
# a_sub_pb2.py
30+
31+
from google.protobuf import descriptor as _descriptor
32+
33+
import a_pb2 as a__pb2
34+
import b_pb2 as b__pb2
35+
36+
from sub.nested import d_pb2 as sub_dot_nested__d__pb2
37+
38+
# ...
39+
```
40+
41+
Using these modules will not work under Python 3, as the imports are not relative. As it can get quite cumbersome to fix these issues, this script will convert the imports automatically.
42+
43+
```bash
44+
fix-protobuf-imports /path/to/python_out/dir
45+
```
46+
47+
This will result in the following working imports:
48+
49+
```python
50+
# a_sub_pb2.py
51+
52+
from google.protobuf import descriptor as _descriptor
53+
54+
from .. import a_pb2 as a__pb2
55+
from .. import b_pb2 as b__pb2
56+
57+
from ..sub.nested import d_pb2 as sub_dot_nested__d__pb2
58+
59+
# ...
60+
```

‎pyproject.toml

+27
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
[build-system]
2+
requires = ["setuptools"]
3+
build-backend = "setuptools.build_meta"
4+
5+
[project]
6+
name = "fix_protobuf_imports"
7+
version = "0.1.4"
8+
authors = [{ name = "Markus Wegmann", email = "[email protected]" }, { name = "Noah Hüsser", email = "[email protected]" }]
9+
description = "A script to fix relative imports (from and to nested sub-directories) within compiled `*_pb2.py` Protobuf files."
10+
readme = "README.md"
11+
license = { file="LICENSE" }
12+
requires-python = ">=3.6"
13+
classifiers = [
14+
"Programming Language :: Python :: 3",
15+
"License :: OSI Approved :: MIT License",
16+
"Operating System :: OS Independent",
17+
]
18+
dependencies = [
19+
"click",
20+
]
21+
22+
[project.urls]
23+
"Homepage" = "https://github.com/technokrat/fix_protobuf_imports"
24+
"Bug Tracker" = "https://github.com/technokrat/fix_protobuf_imports/issues"
25+
26+
[project.scripts]
27+
fix-protobuf-imports = "fix_protobuf_imports.fix_protobuf_imports:main"

‎src/fix_protobuf_imports/__init__.py

Whitespace-only changes.

‎src/fix_protobuf_imports/fix_protobuf_imports.py

+144
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
#!/usr/bin/python3
2+
3+
import os
4+
import re
5+
from pathlib import Path
6+
from typing import Tuple, TypedDict
7+
8+
import click
9+
10+
11+
class ProtobufFilePathInfo(TypedDict):
12+
dir: Path
13+
path: Path
14+
rel_path: Path
15+
16+
17+
@click.command()
18+
@click.option("--dry", is_flag=True, show_default=True, default=False, help="Do not write out the changes to the files.")
19+
@click.argument("root_dir", type=click.Path(exists=True))
20+
def fix_protobuf_imports(root_dir, dry):
21+
"""
22+
A script to fix relative imports (from and to nested sub-directories) within compiled `*_pb2.py` Protobuf files.
23+
"""
24+
25+
root_dir = Path(root_dir)
26+
27+
def generate_lookup(path: Path) -> Tuple[str, ProtobufFilePathInfo]:
28+
name = path.name.split(".")[0]
29+
rel_path = path.relative_to(root_dir)
30+
directory = path.parent.relative_to(root_dir)
31+
32+
return (name, {"dir": directory, "path": path, "rel_path": rel_path})
33+
34+
py_files = list(root_dir.glob("**/*_pb2.py"))
35+
pyi_files = list(root_dir.glob("**/*_pb2.pyi"))
36+
37+
py_files_dictionary = {}
38+
for path in py_files:
39+
name, info = generate_lookup(path)
40+
py_files_dictionary[name] = info
41+
42+
pyi_files_dictionary = {}
43+
for path in pyi_files:
44+
name, info = generate_lookup(path)
45+
pyi_files_dictionary[name] = info
46+
47+
def fix_protobuf_import_in_line(
48+
original_line, referencing_info: ProtobufFilePathInfo, pyi=False
49+
) -> str:
50+
line = original_line
51+
52+
if pyi:
53+
m = re.search(r"^import\s([^\s\.]*_pb2)$", line)
54+
else:
55+
m = re.search(r"^import\s(\S*_pb2)\sas\s(.*)$", line)
56+
57+
if m is not None:
58+
referenced_name = m.group(1)
59+
if pyi:
60+
referenced_alias = None
61+
else:
62+
referenced_alias = m.group(2)
63+
64+
referenced_directory = py_files_dictionary[referenced_name]["dir"]
65+
relative_path_to_referenced_module = os.path.relpath(
66+
referenced_directory, referencing_info["dir"]
67+
)
68+
69+
uppath_levels = relative_path_to_referenced_module.count("..")
70+
71+
original_line = line.replace("\n", "")
72+
73+
downpath = (
74+
relative_path_to_referenced_module.split("..")[-1]
75+
.replace("/", ".")
76+
.replace("\\", ".")
77+
)
78+
if referenced_alias:
79+
line = f'from .{"." * uppath_levels}{downpath if downpath != "." else ""} import {referenced_name} as {referenced_alias}\n'
80+
else:
81+
line = f'from .{"." * uppath_levels}{downpath if downpath != "." else ""} import {referenced_name}\n'
82+
83+
new_line = line.replace("\n", "")
84+
85+
print(f'{referencing_info["rel_path"]}: "{original_line}" -> "{new_line}"')
86+
else:
87+
m = re.search(r"^from\s([^\s\.]+[\S]*)\simport\s(.*_pb2)$", line)
88+
89+
if m is not None:
90+
import_path = m.group(1).replace(".", "/")
91+
92+
referenced_directory = root_dir / import_path
93+
94+
if referenced_directory.exists():
95+
relative_path_to_root = os.path.relpath(
96+
root_dir, referencing_info["dir"]
97+
)
98+
99+
uppath_levels = relative_path_to_root.count("..")
100+
101+
original_line = line.replace("\n", "")
102+
103+
line = (
104+
f'from .{"." * uppath_levels}{m.group(1)} import {m.group(2)}\n'
105+
)
106+
107+
new_line = line.replace("\n", "")
108+
109+
print(
110+
f'{referencing_info["rel_path"]}: "{original_line}" -> "{new_line}"'
111+
)
112+
113+
return line
114+
115+
def fix_protobuf_imports_in_file(name, info: ProtobufFilePathInfo, pyi=False):
116+
with open(info["path"], "r+" if not dry else "r") as f:
117+
lines = f.readlines()
118+
if not dry:
119+
f.seek(0)
120+
121+
for line in lines:
122+
line = fix_protobuf_import_in_line(line, info, pyi)
123+
124+
if not dry:
125+
f.writelines([line])
126+
127+
if not dry:
128+
f.truncate()
129+
f.close()
130+
131+
for (name, info) in py_files_dictionary.items():
132+
fix_protobuf_imports_in_file(name, info)
133+
134+
for (
135+
name,
136+
info,
137+
) in pyi_files_dictionary.items():
138+
fix_protobuf_imports_in_file(name, info, pyi=True)
139+
140+
def main():
141+
fix_protobuf_imports()
142+
143+
if __name__ == "__main__":
144+
main()

0 commit comments

Comments
 (0)