cpython/Objects/unionobject.c
Yurii Karabas 8182c8329c
bpo-44662: Add ability to annotate types.Union (#27214)
Co-authored-by: Ken Jin <28750310+Fidget-Spinner@users.noreply.github.com>
2021-07-29 21:44:48 +02:00

497 lines
13 KiB
C

// types.UnionType -- used to represent e.g. Union[int, str], int | str
#include "Python.h"
#include "pycore_object.h" // _PyObject_GC_TRACK/UNTRACK
#include "pycore_unionobject.h"
#include "structmember.h"
static PyObject *make_union(PyObject *);
typedef struct {
PyObject_HEAD
PyObject *args;
PyObject *parameters;
} unionobject;
static void
unionobject_dealloc(PyObject *self)
{
unionobject *alias = (unionobject *)self;
_PyObject_GC_UNTRACK(self);
Py_XDECREF(alias->args);
Py_XDECREF(alias->parameters);
Py_TYPE(self)->tp_free(self);
}
static int
union_traverse(PyObject *self, visitproc visit, void *arg)
{
unionobject *alias = (unionobject *)self;
Py_VISIT(alias->args);
Py_VISIT(alias->parameters);
return 0;
}
static Py_hash_t
union_hash(PyObject *self)
{
unionobject *alias = (unionobject *)self;
PyObject *args = PyFrozenSet_New(alias->args);
if (args == NULL) {
return (Py_hash_t)-1;
}
Py_hash_t hash = PyObject_Hash(args);
Py_DECREF(args);
return hash;
}
static int
is_generic_alias_in_args(PyObject *args)
{
Py_ssize_t nargs = PyTuple_GET_SIZE(args);
for (Py_ssize_t iarg = 0; iarg < nargs; iarg++) {
PyObject *arg = PyTuple_GET_ITEM(args, iarg);
if (_PyGenericAlias_Check(arg)) {
return 0;
}
}
return 1;
}
static PyObject *
union_instancecheck(PyObject *self, PyObject *instance)
{
unionobject *alias = (unionobject *) self;
Py_ssize_t nargs = PyTuple_GET_SIZE(alias->args);
if (!is_generic_alias_in_args(alias->args)) {
PyErr_SetString(PyExc_TypeError,
"isinstance() argument 2 cannot contain a parameterized generic");
return NULL;
}
for (Py_ssize_t iarg = 0; iarg < nargs; iarg++) {
PyObject *arg = PyTuple_GET_ITEM(alias->args, iarg);
if (PyType_Check(arg)) {
int res = PyObject_IsInstance(instance, arg);
if (res < 0) {
return NULL;
}
if (res) {
Py_RETURN_TRUE;
}
}
}
Py_RETURN_FALSE;
}
static PyObject *
union_subclasscheck(PyObject *self, PyObject *instance)
{
if (!PyType_Check(instance)) {
PyErr_SetString(PyExc_TypeError, "issubclass() arg 1 must be a class");
return NULL;
}
unionobject *alias = (unionobject *)self;
if (!is_generic_alias_in_args(alias->args)) {
PyErr_SetString(PyExc_TypeError,
"issubclass() argument 2 cannot contain a parameterized generic");
return NULL;
}
Py_ssize_t nargs = PyTuple_GET_SIZE(alias->args);
for (Py_ssize_t iarg = 0; iarg < nargs; iarg++) {
PyObject *arg = PyTuple_GET_ITEM(alias->args, iarg);
if (PyType_Check(arg)) {
int res = PyObject_IsSubclass(instance, arg);
if (res < 0) {
return NULL;
}
if (res) {
Py_RETURN_TRUE;
}
}
}
Py_RETURN_FALSE;
}
static PyObject *
union_richcompare(PyObject *a, PyObject *b, int op)
{
if (!_PyUnion_Check(b) || (op != Py_EQ && op != Py_NE)) {
Py_RETURN_NOTIMPLEMENTED;
}
PyObject *a_set = PySet_New(((unionobject*)a)->args);
if (a_set == NULL) {
return NULL;
}
PyObject *b_set = PySet_New(((unionobject*)b)->args);
if (b_set == NULL) {
Py_DECREF(a_set);
return NULL;
}
PyObject *result = PyObject_RichCompare(a_set, b_set, op);
Py_DECREF(b_set);
Py_DECREF(a_set);
return result;
}
static PyObject*
flatten_args(PyObject* args)
{
Py_ssize_t arg_length = PyTuple_GET_SIZE(args);
Py_ssize_t total_args = 0;
// Get number of total args once it's flattened.
for (Py_ssize_t i = 0; i < arg_length; i++) {
PyObject *arg = PyTuple_GET_ITEM(args, i);
if (_PyUnion_Check(arg)) {
total_args += PyTuple_GET_SIZE(((unionobject*) arg)->args);
} else {
total_args++;
}
}
// Create new tuple of flattened args.
PyObject *flattened_args = PyTuple_New(total_args);
if (flattened_args == NULL) {
return NULL;
}
Py_ssize_t pos = 0;
for (Py_ssize_t i = 0; i < arg_length; i++) {
PyObject *arg = PyTuple_GET_ITEM(args, i);
if (_PyUnion_Check(arg)) {
PyObject* nested_args = ((unionobject*)arg)->args;
Py_ssize_t nested_arg_length = PyTuple_GET_SIZE(nested_args);
for (Py_ssize_t j = 0; j < nested_arg_length; j++) {
PyObject* nested_arg = PyTuple_GET_ITEM(nested_args, j);
Py_INCREF(nested_arg);
PyTuple_SET_ITEM(flattened_args, pos, nested_arg);
pos++;
}
} else {
if (arg == Py_None) {
arg = (PyObject *)&_PyNone_Type;
}
Py_INCREF(arg);
PyTuple_SET_ITEM(flattened_args, pos, arg);
pos++;
}
}
assert(pos == total_args);
return flattened_args;
}
static PyObject*
dedup_and_flatten_args(PyObject* args)
{
args = flatten_args(args);
if (args == NULL) {
return NULL;
}
Py_ssize_t arg_length = PyTuple_GET_SIZE(args);
PyObject *new_args = PyTuple_New(arg_length);
if (new_args == NULL) {
Py_DECREF(args);
return NULL;
}
// Add unique elements to an array.
Py_ssize_t added_items = 0;
for (Py_ssize_t i = 0; i < arg_length; i++) {
int is_duplicate = 0;
PyObject* i_element = PyTuple_GET_ITEM(args, i);
for (Py_ssize_t j = 0; j < added_items; j++) {
PyObject* j_element = PyTuple_GET_ITEM(new_args, j);
int is_ga = _PyGenericAlias_Check(i_element) &&
_PyGenericAlias_Check(j_element);
// RichCompare to also deduplicate GenericAlias types (slower)
is_duplicate = is_ga ? PyObject_RichCompareBool(i_element, j_element, Py_EQ)
: i_element == j_element;
// Should only happen if RichCompare fails
if (is_duplicate < 0) {
Py_DECREF(args);
Py_DECREF(new_args);
return NULL;
}
if (is_duplicate)
break;
}
if (!is_duplicate) {
Py_INCREF(i_element);
PyTuple_SET_ITEM(new_args, added_items, i_element);
added_items++;
}
}
Py_DECREF(args);
_PyTuple_Resize(&new_args, added_items);
return new_args;
}
static int
is_unionable(PyObject *obj)
{
return (obj == Py_None ||
PyType_Check(obj) ||
_PyGenericAlias_Check(obj) ||
_PyUnion_Check(obj));
}
PyObject *
_Py_union_type_or(PyObject* self, PyObject* other)
{
if (!is_unionable(self) || !is_unionable(other)) {
Py_RETURN_NOTIMPLEMENTED;
}
PyObject *tuple = PyTuple_Pack(2, self, other);
if (tuple == NULL) {
return NULL;
}
PyObject *new_union = make_union(tuple);
Py_DECREF(tuple);
return new_union;
}
static int
union_repr_item(_PyUnicodeWriter *writer, PyObject *p)
{
_Py_IDENTIFIER(__module__);
_Py_IDENTIFIER(__qualname__);
_Py_IDENTIFIER(__origin__);
_Py_IDENTIFIER(__args__);
PyObject *qualname = NULL;
PyObject *module = NULL;
PyObject *tmp;
PyObject *r = NULL;
int err;
if (p == (PyObject *)&_PyNone_Type) {
return _PyUnicodeWriter_WriteASCIIString(writer, "None", 4);
}
if (_PyObject_LookupAttrId(p, &PyId___origin__, &tmp) < 0) {
goto exit;
}
if (tmp) {
Py_DECREF(tmp);
if (_PyObject_LookupAttrId(p, &PyId___args__, &tmp) < 0) {
goto exit;
}
if (tmp) {
// It looks like a GenericAlias
Py_DECREF(tmp);
goto use_repr;
}
}
if (_PyObject_LookupAttrId(p, &PyId___qualname__, &qualname) < 0) {
goto exit;
}
if (qualname == NULL) {
goto use_repr;
}
if (_PyObject_LookupAttrId(p, &PyId___module__, &module) < 0) {
goto exit;
}
if (module == NULL || module == Py_None) {
goto use_repr;
}
// Looks like a class
if (PyUnicode_Check(module) &&
_PyUnicode_EqualToASCIIString(module, "builtins"))
{
// builtins don't need a module name
r = PyObject_Str(qualname);
goto exit;
}
else {
r = PyUnicode_FromFormat("%S.%S", module, qualname);
goto exit;
}
use_repr:
r = PyObject_Repr(p);
exit:
Py_XDECREF(qualname);
Py_XDECREF(module);
if (r == NULL) {
return -1;
}
err = _PyUnicodeWriter_WriteStr(writer, r);
Py_DECREF(r);
return err;
}
static PyObject *
union_repr(PyObject *self)
{
unionobject *alias = (unionobject *)self;
Py_ssize_t len = PyTuple_GET_SIZE(alias->args);
_PyUnicodeWriter writer;
_PyUnicodeWriter_Init(&writer);
for (Py_ssize_t i = 0; i < len; i++) {
if (i > 0 && _PyUnicodeWriter_WriteASCIIString(&writer, " | ", 3) < 0) {
goto error;
}
PyObject *p = PyTuple_GET_ITEM(alias->args, i);
if (union_repr_item(&writer, p) < 0) {
goto error;
}
}
return _PyUnicodeWriter_Finish(&writer);
error:
_PyUnicodeWriter_Dealloc(&writer);
return NULL;
}
static PyMemberDef union_members[] = {
{"__args__", T_OBJECT, offsetof(unionobject, args), READONLY},
{0}
};
static PyMethodDef union_methods[] = {
{"__instancecheck__", union_instancecheck, METH_O},
{"__subclasscheck__", union_subclasscheck, METH_O},
{0}};
static PyObject *
union_getitem(PyObject *self, PyObject *item)
{
unionobject *alias = (unionobject *)self;
// Populate __parameters__ if needed.
if (alias->parameters == NULL) {
alias->parameters = _Py_make_parameters(alias->args);
if (alias->parameters == NULL) {
return NULL;
}
}
PyObject *newargs = _Py_subs_parameters(self, alias->args, alias->parameters, item);
if (newargs == NULL) {
return NULL;
}
PyObject *res;
Py_ssize_t nargs = PyTuple_GET_SIZE(newargs);
if (nargs == 0) {
res = make_union(newargs);
}
else {
res = PyTuple_GET_ITEM(newargs, 0);
Py_INCREF(res);
for (Py_ssize_t iarg = 1; iarg < nargs; iarg++) {
PyObject *arg = PyTuple_GET_ITEM(newargs, iarg);
Py_SETREF(res, PyNumber_Or(res, arg));
if (res == NULL) {
break;
}
}
}
Py_DECREF(newargs);
return res;
}
static PyMappingMethods union_as_mapping = {
.mp_subscript = union_getitem,
};
static PyObject *
union_parameters(PyObject *self, void *Py_UNUSED(unused))
{
unionobject *alias = (unionobject *)self;
if (alias->parameters == NULL) {
alias->parameters = _Py_make_parameters(alias->args);
if (alias->parameters == NULL) {
return NULL;
}
}
Py_INCREF(alias->parameters);
return alias->parameters;
}
static PyGetSetDef union_properties[] = {
{"__parameters__", union_parameters, (setter)NULL, "Type variables in the types.UnionType.", NULL},
{0}
};
static PyNumberMethods union_as_number = {
.nb_or = _Py_union_type_or, // Add __or__ function
};
static const char* const cls_attrs[] = {
"__module__", // Required for compatibility with typing module
NULL,
};
static PyObject *
union_getattro(PyObject *self, PyObject *name)
{
unionobject *alias = (unionobject *)self;
if (PyUnicode_Check(name)) {
for (const char * const *p = cls_attrs; ; p++) {
if (*p == NULL) {
break;
}
if (_PyUnicode_EqualToASCIIString(name, *p)) {
return PyObject_GetAttr((PyObject *) Py_TYPE(alias), name);
}
}
}
return PyObject_GenericGetAttr(self, name);
}
PyTypeObject _PyUnion_Type = {
PyVarObject_HEAD_INIT(&PyType_Type, 0)
.tp_name = "types.UnionType",
.tp_doc = "Represent a PEP 604 union type\n"
"\n"
"E.g. for int | str",
.tp_basicsize = sizeof(unionobject),
.tp_dealloc = unionobject_dealloc,
.tp_alloc = PyType_GenericAlloc,
.tp_free = PyObject_GC_Del,
.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
.tp_traverse = union_traverse,
.tp_hash = union_hash,
.tp_getattro = union_getattro,
.tp_members = union_members,
.tp_methods = union_methods,
.tp_richcompare = union_richcompare,
.tp_as_mapping = &union_as_mapping,
.tp_as_number = &union_as_number,
.tp_repr = union_repr,
.tp_getset = union_properties,
};
static PyObject *
make_union(PyObject *args)
{
assert(PyTuple_CheckExact(args));
args = dedup_and_flatten_args(args);
if (args == NULL) {
return NULL;
}
if (PyTuple_GET_SIZE(args) == 1) {
PyObject *result1 = PyTuple_GET_ITEM(args, 0);
Py_INCREF(result1);
Py_DECREF(args);
return result1;
}
unionobject *result = PyObject_GC_New(unionobject, &_PyUnion_Type);
if (result == NULL) {
Py_DECREF(args);
return NULL;
}
result->parameters = NULL;
result->args = args;
_PyObject_GC_TRACK(result);
return (PyObject*)result;
}