Repository: mandiant/capa
File: capa/features/freeze/__init__.py (Python)

Code before:
import capa.features.common
import capa.features.address
import capa.features.basicblock
import capa.features.extractors.base_extractor
from capa.helpers import assert_never
from capa.features.freeze.features import Feature, feature_from_capa

logger = logging.getLogger(__name__)


class HashableModel(BaseModel):
    model_config = ConfigDict(frozen=True)


class AddressType(str, Enum):
    ABSOLUTE = "absolute"
    RELATIVE = "relative"
    FILE = "file"
    DN_TOKEN = "dn token"
    DN_TOKEN_OFFSET = "dn token offset"
    NO_ADDRESS = "no address"


class Address(HashableModel):
    type: AddressType
    value: Union[int, Tuple[int, int], None] = None

    @classmethod
    def from_capa(cls, a: capa.features.address.Address) -> "Address":
        if isinstance(a, capa.features.address.AbsoluteVirtualAddress):
            return cls(type=AddressType.ABSOLUTE, value=int(a))

        elif isinstance(a, capa.features.address.RelativeVirtualAddress):
            return cls(type=AddressType.RELATIVE, value=int(a))

        elif isinstance(a, capa.features.address.FileOffsetAddress):
            return cls(type=AddressType.FILE, value=int(a))

        elif isinstance(a, capa.features.address.DNTokenAddress):
            return cls(type=AddressType.DN_TOKEN, value=int(a))

        elif isinstance(a, capa.features.address.DNTokenOffsetAddress):
            return cls(type=AddressType.DN_TOKEN_OFFSET, value=(a.token, a.offset))

        elif a == capa.features.address.NO_ADDRESS or isinstance(a, capa.features.address._NoAddress):
            return cls(type=AddressType.NO_ADDRESS, value=None)

        elif isinstance(a, capa.features.address.Address) and not issubclass(type(a), capa.features.address.Address):
            raise ValueError("don't use an Address instance directly")

        elif isinstance(a, capa.features.address.Address):

Review comment (suggestion):
```suggestion
    value: Union[int, Tuple[int, int], None] = None  # None default value to support deserialization of NO_ADDRESS
```

Code after:
import capa.features.common
import capa.features.address
import capa.features.basicblock
import capa.features.extractors.base_extractor
from capa.helpers import assert_never
from capa.features.freeze.features import Feature, feature_from_capa

logger = logging.getLogger(__name__)


class HashableModel(BaseModel):
    model_config = ConfigDict(frozen=True)


class AddressType(str, Enum):
    ABSOLUTE = "absolute"
    RELATIVE = "relative"
    FILE = "file"
    DN_TOKEN = "dn token"
    DN_TOKEN_OFFSET = "dn token offset"
    NO_ADDRESS = "no address"


class Address(HashableModel):
    type: AddressType
    value: Union[int, Tuple[int, int], None] = None  # None default value to support deserialization of NO_ADDRESS

    @classmethod
    def from_capa(cls, a: capa.features.address.Address) -> "Address":
        if isinstance(a, capa.features.address.AbsoluteVirtualAddress):
            return cls(type=AddressType.ABSOLUTE, value=int(a))

        elif isinstance(a, capa.features.address.RelativeVirtualAddress):
            return cls(type=AddressType.RELATIVE, value=int(a))

        elif isinstance(a, capa.features.address.FileOffsetAddress):
            return cls(type=AddressType.FILE, value=int(a))

        elif isinstance(a, capa.features.address.DNTokenAddress):
            return cls(type=AddressType.DN_TOKEN, value=int(a))

        elif isinstance(a, capa.features.address.DNTokenOffsetAddress):
            return cls(type=AddressType.DN_TOKEN_OFFSET, value=(a.token, a.offset))

        elif a == capa.features.address.NO_ADDRESS or isinstance(a, capa.features.address._NoAddress):
            return cls(type=AddressType.NO_ADDRESS, value=None)

        elif isinstance(a, capa.features.address.Address) and not issubclass(type(a), capa.features.address.Address):
            raise ValueError("don't use an Address instance directly")

        elif isinstance(a, capa.features.address.Address):