在继承的数据类中使用 __new__

Question

假设我有以下代码用于处理个人和国家之间的链接：

from dataclasses import dataclass

@dataclass
class Country:
    iso2 : str
    iso3 : str
    name : str

countries = [ Country('AW','ABW','Aruba'),
              Country('AF','AFG','Afghanistan'),
              Country('AO','AGO','Angola')]
countries_by_iso2 = {c.iso2 : c for c in countries}
countries_by_iso3 = {c.iso3 : c for c in countries}

@dataclass
class CountryLink:
    person_id : int
    country : Country

country_links = [ CountryLink(123, countries_by_iso2['AW']),
                  CountryLink(456, countries_by_iso3['AFG']),
                  CountryLink(789, countries_by_iso2['AO'])]

print(country_links[0].country.name)

这一切正常运作，但我决定让它更加流畅，以处理不同形式的输入。我还想使用__new__来确保每次都获得有效的ISO代码，并且希望在这种情况下对象创建失败。因此，我添加了几个继承自此类的新类：

@dataclass
class CountryLinkFromISO2(CountryLink):
    def __new__(cls, person_id : int, iso2 : str):
        if iso2 not in countries_by_iso2:
            return None
        new_obj = super().__new__(cls)
        new_obj.country = countries_by_iso2[iso2]
        return new_obj

@dataclass
class CountryLinkFromISO3(CountryLink):
    def __new__(cls, person_id : int, iso3 : str):
        if iso3 not in countries_by_iso3:
            return None
        new_obj = super().__new__(cls)
        new_obj.country = countries_by_iso3[iso3]
        return new_obj

country_links = [ CountryLinkFromISO2(123, 'AW'),
                  CountryLinkFromISO3(456, 'AFG'),
                  CountryLinkFromISO2(789, 'AO')]

这看起来乍一看是可以的，但是我遇到了一个问题：

a = CountryLinkFromISO2(123, 'AW')
print(type(a))
print(a.country)
print(type(a.country))

返回：

<class '__main__.CountryLinkFromISO2'>
AW
<class 'str'>

继承的对象具有正确的类型，但是它的属性country只是一个字符串，而不是我所期望的Country类型。在__new__中，我放置了打印语句来检查new_obj.country的类型，在return行之前是正确的。

我想要实现的是让a成为CountryLinkFromISO2类型的对象，继承我对CountryLink所做的更改，并且它具有从字典countries_by_iso2中获取的属性country。我该如何实现这一点？

- EdG

---

你确定要覆盖__new__而不是__init__吗？你也可以考虑使用像attrs这样的库。 - Nathaniel Ford

@NathanielFord 如果输入无效，__init__ 将始终返回类的实例，这不是我想要发生的。我可以有一个引发异常的 __init__，但这意味着每次尝试调用我的代码时都必须将其放入 try/except 块中，这很笨拙并可能导致性能问题。 - EdG

我认为Mark提供了正确的课程（工厂方法），但你应该查看attrs或类似库的验证器。 - Nathaniel Ford

Answer 1

即使数据类在幕后自动完成了这一过程，也不意味着您的类没有 __init__()。它们有，并且看起来像：

def __init__(self, person_id: int, country: Country):
    self.person_id = person_id
    self.country = country

当你使用以下代码创建类时：

CountryLinkFromISO2(123, 'AW')

这个"AW"字符串被传递到__init__()中，并将该值设置为一个字符串。

以这种方式使用__new__()是脆弱的，从构造函数返回None相当不符合Python风格（在我看来）。也许您最好制作一个实际的工厂函数，它返回None或您想要的类。然后您就不需要完全搞清楚__new__()了。

@dataclass
class CountryLinkFromISO2(CountryLink):
    @classmethod
    def from_country_code(cls, person_id : int, iso2 : str):
        if iso2 not in countries_by_iso2:
            return None
        return cls(person_id, countries_by_iso2[iso2])

a = CountryLinkFromISO2.from_country_code(123, 'AW')

如果由于某些原因它需要与 __new__() 一起工作，那么当没有匹配时，您可以从 new 返回 None ，并在 __post_init__() 中设置国家：

@dataclass
class CountryLinkFromISO2(CountryLink):
    def __new__(cls, person_id : int, iso2 : str):
        if iso2 not in countries_by_iso2:
            return None
        return super().__new__(cls)
    
    def __post_init__(self):        
        self.country = countries_by_iso2[self.country]

- Mark

---

这是否意味着由dataclass装饰器隐式创建的__init__在具有相同位置参数的__new__之后运行，因此覆盖了先前的country值？从构造函数返回None可能不符合Pythonic风格，但最好的替代方案是使用try/except块，我理解这可能会导致性能问题。 - EdG

是的，__init__() 是由 dataclass 装饰器创建的。它在这里有文档记录，并且是使用 dataclass 的原因之一。而 __init__() 在 __new__() 之后被调用。这不是特定于 dataclass 的。 - Mark

那么，如果我想避免__init__替换在__new__中设置的country值，我需要手动指定__init__和__new__，使它们使用相同的位置参数但故意不覆盖吗？ - EdG

Answer 2

您看到的行为是由于数据类在__init__中设置它们的字段，而此时__new__已经执行。

解决这个问题的Pythonic方法是提供一个替代构造函数。我不会创建子类，因为它们仅用于其构造函数。

例如：

@dataclass
class CountryLink:
    person_id: int
    country: Country

    @classmethod
    def from_iso2(cls, person_id: int, country_code: str):
        try:
            return cls(person_id, countries_by_iso2[country_code])
        except KeyError:
            raise ValueError(f'invalid ISO2 country code {country_code!r}') from None

    @classmethod
    def from_iso3(cls, person_id: int, country_code: str):
        try:
            return cls(person_id, countries_by_iso3[country_code])
        except KeyError:
            raise ValueError(f'invalid ISO3 country code {country_code!r}') from None

country_links = [ CountryLink.from_iso2(123, 'AW'),
                  CountryLink.from_iso3(456, 'AFG'),
                  CountryLink.from_iso2(789, 'AO')]

- Jasmijn

---

这种方法意味着每次创建类时都必须使用try/except块，这会对性能产生影响吗？ - EdG

如果没有引发异常，try/except块的性能影响很小。甚至可能会更快，因为没有'in'检查和if分支。无论如何，我认为这不太可能成为您代码的瓶颈。 - Jasmijn

有人可能认为强制检查无效数据是一件好事。 - Matthew Purdon