[Python] How do I add missing imports to a Python module using ast-grep?

I have the following ast-grep rule to replace some http status integer literals in my code with constants

id: replace_http_status_with_constants
language: python
rule:
all:
- regex: "^(200|201|204|302|400|401|403|404|500|502|503|504)$"
- any:
- pattern: "$STATUS"
kind: integer
inside:
kind: "comparison_operator"
- pattern: "$STATUS"
kind: integer
inside:
kind: "keyword_argument"
has:
kind: identifier
regex: "status.*"

- pattern: "$STATUS"
kind: integer
inside:
kind: "assignment"
has:
stopBy: end
kind: identifier
regex: "status*"

- pattern: "$STATUS"
kind: "integer"
inside:
kind: "decorator"
stopBy:
kind: decorator
has:
kind: "identifier"
regex: "pytest|extend_schema"
stopBy: end

rewriters:
- id: HTTP_200_OK
rule:
regex: "200"
fix: "HTTP_200_OK"
- id: HTTP_201_CREATED
rule:
regex: "201"
fix: "HTTP_201_CREATED"
- id: HTTP_204_NO_CONTENT
rule:
regex: "204"
fix: "HTTP_204_NO_CONTENT"
- id: HTTP_302_FOUND
rule:
regex: "302"
fix: "HTTP_302_FOUND"
- id: HTTP_400_BAD_REQUEST
rule:
regex: "400"
fix: "HTTP_400_BAD_REQUEST"
- id: HTTP_401_UNAUTHORIZED
rule:
regex: "401"
fix: "HTTP_401_UNAUTHORIZED"
- id: HTTP_403_FORBIDDEN
rule:
regex: "403"
fix: "HTTP_403_FORBIDDEN"
- id: HTTP_404_NOT_FOUND
rule:
regex: "404"
fix: "HTTP_404_NOT_FOUND"
- id: HTTP_500_INTERNAL_SERVER_ERROR
rule:
regex: "500"
fix: "HTTP_500_INTERNAL_SERVER_ERROR"
- id: HTTP_502_BAD_GATEWAY
rule:
regex: "502"
fix: "HTTP_502_BAD_GATEWAY"
- id: HTTP_503_SERVICE_UNAVAILABLE
rule:
regex: "503"
fix: "HTTP_503_SERVICE_UNAVAILABLE"
- id: HTTP_504_GATEWAY_TIMEOUT
rule:
regex: "504"
fix: "HTTP_504_GATEWAY_TIMEOUT"

transform:
NEW_STATUS:
rewrite:
source: "$STATUS"
rewriters:
- HTTP_200_OK
- HTTP_201_CREATED
- HTTP_204_NO_CONTENT
- HTTP_302_FOUND
- HTTP_400_BAD_REQUEST
- HTTP_401_UNAUTHORIZED
- HTTP_403_FORBIDDEN
- HTTP_404_NOT_FOUND
- HTTP_500_INTERNAL_SERVER_ERROR
- HTTP_502_BAD_GATEWAY
- HTTP_503_SERVICE_UNAVAILABLE
- HTTP_504_GATEWAY_TIMEOUT

fix: $NEW_STATUS


Is there a way to make it so I can also create any missing imports for these constants?

So far I've come up with this, but it can't really include only missing imports

id: ensure_http_status_imports
language: python
rule:
kind: "module"
pattern: $MODULE
all:
- regex: "HTTP_200_OK|HTTP_201_CREATED|HTTP_204_NO_CONTENT|HTTP_302_FOUND|HTTP_400_BAD_REQUEST|HTTP_401_UNAUTHORIZED|HTTP_403_FORBIDDEN|HTTP_404_NOT_FOUND|HTTP_500_INTERNAL_SERVER_ERROR|HTTP_502_BAD_GATEWAY|HTTP_503_SERVICE_UNAVAILABLE|HTTP_504_GATEWAY_TIMEOUT"
- has:
pattern: $IDENTIFIER
stopBy: end
kind: identifier
not:
inside:
kind: attribute

- not:
has:
pattern: $IDENTIFIER
stopBy: end
kind: "import_from_statement"
has:
kind: identifier
stopBy: end

fix: |
from rest_framework.status import (
HTTP_200_OK,
HTTP_201_CREATED,
HTTP_204_NO_CONTENT,
HTTP_302_FOUND,
HTTP_400_BAD_REQUEST,
HTTP_401_UNAUTHORIZED,
HTTP_403_FORBIDDEN,
HTTP_404_NOT_FOUND,
HTTP_500_INTERNAL_SERVER_ERROR,
HTTP_502_BAD_GATEWAY,
HTTP_503_SERVICE_UNAVAILABLE,
HTTP_504_GATEWAY_TIMEOUT,
)
$MODULE

Continue reading...