Use string value for argument typed as Literal

user_2956276 0
Viewed 225

I use kms.decrypt() method from boto3 package. For typing support I use the boto3-stubs package.

The decrypt method has attribute EncryptionAlgorithm, which is typed as

EncryptionAlgorithmSpecType = Literal["RSAES_OAEP_SHA_1", "RSAES_OAEP_SHA_256", "SYMMETRIC_DEFAULT"]

I parse the encryption algorithm via regular expression from input string so in my case the EncryptionAlgorithm value is str, not a Literal. Mypy complaint for it.

I do not want to disable type check for this line and the only solution I was found is the helper method which convert the str value to literal in this way:

import boto3
from mypy_boto3_kms.literals import EncryptionAlgorithmSpecType

def getEncryptionAlgorithmLiteral(algorithm: str) -> EncryptionAlgorithmSpecType:
  result: EncryptionAlgorithmSpecType
  if algorithm == "RSAES_OAEP_SHA_1":
    result = "RSAES_OAEP_SHA_1"
  elif algorithm == "RSAES_OAEP_SHA_256":
    result = "RSAES_OAEP_SHA_256"
  elif algorithm == "SYMMETRIC_DEFAULT":
    result = "SYMMETRIC_DEFAULT"
  else:
    raise Exception(f"Unexpected algorithm '{algorithm}'. It must be one of {EncryptionAlgorithmSpecType}")
  return result

def main(binaryEncData: bytes, keyId: str, algorithm: str):
  kms = boto3.client('kms')
  kmsResult = kms.decrypt(CiphertextBlob=binaryEncData, 
                          KeyId=keyId,
                          EncryptionAlgorithm=getEncryptionAlgorithmLiteral(algorithm))

I wonder if there is some better way how to achieve the conversion in getEncryptionAlgorithmLiteral() method so I do not need to type all these values twice. Ideally I would like to use values directly from EncryptionAlgorithmSpecType type instead of typing them again to my code.

python-typing mypy python
1 Answers

You can use typing.get_args to get the arguments passed in to typing.Literal. In this case, you'll need to combine it with typing.cast so you can signal to "mypy" that the string value that the function returns is an acceptable Literal value.

from typing import cast
# Import below from `typing` in Python 3.8+
from typing_extensions import Literal, get_args


# noinspection SpellCheckingInspection
EncryptionAlgorithmSpecType = Literal["RSAES_OAEP_SHA_1",
                                      "RSAES_OAEP_SHA_256",
                                      "SYMMETRIC_DEFAULT"]

_valid_algorithms = get_args(EncryptionAlgorithmSpecType)
print(_valid_algorithms)


def get_encryption_algorithm_literal(
        algorithm: str) -> EncryptionAlgorithmSpecType:

    if algorithm not in _valid_algorithms:
        valid_values = str(list(_valid_algorithms)).replace("'", "")
        raise Exception(f"Unexpected algorithm '{algorithm}'. "
                        f"It must be one of {valid_values}")

    # cast string to literal, so static type checkers such as 'mypy'
    # don't complain.
    return cast(EncryptionAlgorithmSpecType, algorithm)


def main():
    # noinspection SpellCheckingInspection
    string = 'RSAES_OAEP_SHA_256'

    my_algorithm = get_encryption_algorithm_literal(string)
    print(type(my_algorithm), my_algorithm)


if __name__ == '__main__':
    main()

Output:

('RSAES_OAEP_SHA_1', 'RSAES_OAEP_SHA_256', 'SYMMETRIC_DEFAULT')
<class 'str'> RSAES_OAEP_SHA_256

Result for an invalid input like 'RSAES_OAEP_SHA_2567':

Traceback (most recent call last):
  ...
    raise Exception(...)
Exception: Unexpected algorithm 'RSAES_OAEP_SHA_2567'. It must be one of [RSAES_OAEP_SHA_1, RSAES_OAEP_SHA_256, SYMMETRIC_DEFAULT]
user_102375060
Related