-
-
Notifications
You must be signed in to change notification settings - Fork 69
/
_cleanup.py
68 lines (51 loc) · 2.29 KB
/
_cleanup.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
import sys
import pkgutil
from pathlib import Path
from typing_extensions import Protocol, runtime_checkable
PRISMA_INIT_CONTENTS = "__title__ = 'prisma'"
@runtime_checkable
class SourceLoader(Protocol):
def get_filename(self) -> str:
...
def main(*args: str) -> None:
if args:
pkg_name = args[0]
else:
pkg_name = 'prisma'
cleanup(pkg_name=pkg_name)
def cleanup(pkg_name: str = 'prisma') -> None:
"""Remove python files that are auto-generated by Prisma Client Python"""
# pkgutil.get_loader() can return a loader that doesn't have access to
# the packages source location. We should provide an easy to understand error
# for this case even if it is incredibly unlikey (if not impossible)
# to happen in our use case
loader = pkgutil.get_loader(pkg_name) # pyright: ignore[reportDeprecated]
if loader is None:
raise RuntimeError(f'Could not resolve package: {pkg_name}')
if not isinstance(loader, SourceLoader):
raise RuntimeError(f'Received unresolvable import loader: {loader}')
# ensure the package we've been given is actually a Prisma Client Python package
# we don't want to make it easy for users to accidentaly delete files from
# a non Prisma Client Python package
pkg_path = Path(loader.get_filename())
if PRISMA_INIT_CONTENTS not in pkg_path.read_text():
raise RuntimeError(
'The given package does not appear to be a Prisma Client Python package.'
)
# as we rely on prisma to cleanup the templates for us
# we have to make sure that prisma is importable and
# if any template rendered incorrect syntax or any other
# kind of error that wasn't automatically cleaned up,
# prisma will raise an error when imported,
# removing prisma/client.py fixes this as it is
# the only default entrypoint to generated code.
file = pkg_path.parent / 'client.py'
if file.exists():
file.unlink()
# the `prisma` package will always exist even when using a custom output
# location so it is safe to use here
from prisma.generator.generator import cleanup_templates
cleanup_templates(rootdir=pkg_path.parent)
print(f'Successfully removed all auto-generated files from {pkg_path}')
if __name__ == '__main__':
main(*sys.argv[1:])