from newtype import NewType, newtype_exclude from enum import Enum import pytest class ENV(NewType(str), Enum): LOCAL = "LOCAL" DEV = "DEV" SIT = "SIT" UAT = "UAT" PREPROD = "PREPROD" PROD = "PROD" class RegularENV(str, Enum): LOCAL = "LOCAL" DEV = "DEV" SIT = "SIT" UAT = "UAT" PREPROD = "PREPROD" PROD = "PROD" class ENVVariant(str): __VALID_MEMBERS__ = ["LOCAL", "DEV", "SIT", "UAT", "PREPROD", "PROD"] def __new__(cls, value: str): members = ENVVariant.__VALID_MEMBERS__ if isinstance(value, RollYourOwnNewTypeEnum): value_as_str = str(value.value) else: value_as_str = str(value) if value_as_str not in members: raise ValueError(f"`value` = {value} must be one of `{members}`; `value_as_str` = {value_as_str}") return super().__new__(cls, value_as_str) # why not i write my own `.replace(..)` # yes, you can but how? def my_replace(self, old: "ENVVariant", new: "ENVVariant", count: int=-1): return ENVVariant(str(self).replace(str(old), str(new), count)) class RollYourOwnNewTypeEnum(ENVVariant, Enum): LOCAL = "LOCAL" DEV = "DEV" SIT = "SIT" UAT = "UAT" PREPROD = "PREPROD" PROD = "PROD" def test_nt_env_replace(): env = ENV.LOCAL assert env is ENV.LOCAL assert env is not ENV.DEV assert isinstance(env, ENV) # let's say now we want to replace the environment # nevermind about the reason why we want to do so env = env.replace(ENV.LOCAL, ENV.DEV) # replacement is successful assert env is ENV.DEV assert env is not ENV.LOCAL # still an `ENV` assert isinstance(env, ENV) assert isinstance(env, str) with pytest.raises(ValueError): # cannot replace with something that is not a `ENV` env = env.replace(ENV.DEV, "NotAnEnv") with pytest.raises(ValueError): # cannot even make 'DEV' -> 'dev' env = env.lower() def test_reg_env_replace(): env = RegularENV.LOCAL # expected outcomes assert env is RegularENV.LOCAL # pass assert env is not RegularENV.DEV # pass assert isinstance(env, RegularENV) # pass # now we try to replace env = env.replace(RegularENV.LOCAL, RegularENV.DEV) # we are hoping that it will continue to be a `RegularENV.DEV` but it is not assert env is not RegularENV.DEV # pass, no longer a `RegularENV` assert env is not RegularENV.LOCAL # pass, no longer a `RegularENV` assert not isinstance(env, RegularENV) assert isinstance(env, str) # 'downcast' (?) to `str` def test_ryont_env_replace(): env = RollYourOwnNewTypeEnum.LOCAL # expected outcomes assert env is RollYourOwnNewTypeEnum.LOCAL # pass assert env is not RollYourOwnNewTypeEnum.DEV # pass assert isinstance(env, RollYourOwnNewTypeEnum) # pass # now we try to replace env = env.replace(RollYourOwnNewTypeEnum.LOCAL, RollYourOwnNewTypeEnum.DEV) # we are hoping that it will continue to be a `RollYourOwnNewTypeEnum.DEV` but it is not assert env is not RollYourOwnNewTypeEnum.DEV # pass, no longer a `RollYourOwnNewTypeEnum` assert env is not RollYourOwnNewTypeEnum.LOCAL # pass, no longer a `RollYourOwnNewTypeEnum` assert not isinstance(env, RollYourOwnNewTypeEnum) assert isinstance(env, str) # 'downcast' (?) to `str` with pytest.raises(AssertionError): assert env is RollYourOwnNewTypeEnum.DEV with pytest.raises(AssertionError): assert env is RollYourOwnNewTypeEnum.DEV with pytest.raises(AssertionError): assert isinstance(env, RollYourOwnNewTypeEnum) env = env.replace("DEV", "NotAnEnv") assert env == "NotAnEnv" # this 'shouldn't' pass but it does env = RollYourOwnNewTypeEnum.LOCAL env = env.my_replace(RollYourOwnNewTypeEnum.LOCAL, RollYourOwnNewTypeEnum.PREPROD) assert isinstance(env, str) assert env is RollYourOwnNewTypeEnum.PREPROD assert isinstance(env, RollYourOwnNewTypeEnum)