From ba6fdcf6f8c618c580631cfedbcd31eaf104a0c2 Mon Sep 17 00:00:00 2001
From: 97littleleaf11 <97littleleaf11@gmail.com>
Date: Tue, 6 Apr 2021 17:28:34 +0800
Subject: [PATCH 1/6] Move dict_new related ops to the first section of
 dict_ops.py

---
 mypyc/primitives/dict_ops.py | 68 ++++++++++++++++++------------------
 1 file changed, 34 insertions(+), 34 deletions(-)

diff --git a/mypyc/primitives/dict_ops.py b/mypyc/primitives/dict_ops.py
index ff9b6482a782..62a4151c97f6 100644
--- a/mypyc/primitives/dict_ops.py
+++ b/mypyc/primitives/dict_ops.py
@@ -17,6 +17,40 @@
     type=object_rprimitive,
     src='PyDict_Type')
 
+# Construct an empty dictionary.
+dict_new_op = custom_op(
+    arg_types=[],
+    return_type=dict_rprimitive,
+    c_function_name='PyDict_New',
+    error_kind=ERR_MAGIC)
+
+# Construct a dictionary from keys and values.
+# Positional argument is the number of key-value pairs
+# Variable arguments are (key1, value1, ..., keyN, valueN).
+dict_build_op = custom_op(
+    arg_types=[c_pyssize_t_rprimitive],
+    return_type=dict_rprimitive,
+    c_function_name='CPyDict_Build',
+    error_kind=ERR_MAGIC,
+    var_arg_type=object_rprimitive)
+
+# Construct a dictionary from another dictionary.
+function_op(
+    name='builtins.dict',
+    arg_types=[dict_rprimitive],
+    return_type=dict_rprimitive,
+    c_function_name='PyDict_Copy',
+    error_kind=ERR_MAGIC,
+    priority=2)
+
+# Generic one-argument dict constructor: dict(obj)
+function_op(
+    name='builtins.dict',
+    arg_types=[object_rprimitive],
+    return_type=dict_rprimitive,
+    c_function_name='CPyDict_FromAny',
+    error_kind=ERR_MAGIC)
+
 # dict[key]
 dict_get_item_op = method_op(
     name='__getitem__',
@@ -84,40 +118,6 @@
     c_function_name='CPyDict_GetWithNone',
     error_kind=ERR_MAGIC)
 
-# Construct an empty dictionary.
-dict_new_op = custom_op(
-    arg_types=[],
-    return_type=dict_rprimitive,
-    c_function_name='PyDict_New',
-    error_kind=ERR_MAGIC)
-
-# Construct a dictionary from keys and values.
-# Positional argument is the number of key-value pairs
-# Variable arguments are (key1, value1, ..., keyN, valueN).
-dict_build_op = custom_op(
-    arg_types=[c_pyssize_t_rprimitive],
-    return_type=dict_rprimitive,
-    c_function_name='CPyDict_Build',
-    error_kind=ERR_MAGIC,
-    var_arg_type=object_rprimitive)
-
-# Construct a dictionary from another dictionary.
-function_op(
-    name='builtins.dict',
-    arg_types=[dict_rprimitive],
-    return_type=dict_rprimitive,
-    c_function_name='PyDict_Copy',
-    error_kind=ERR_MAGIC,
-    priority=2)
-
-# Generic one-argument dict constructor: dict(obj)
-function_op(
-    name='builtins.dict',
-    arg_types=[object_rprimitive],
-    return_type=dict_rprimitive,
-    c_function_name='CPyDict_FromAny',
-    error_kind=ERR_MAGIC)
-
 # dict.keys()
 method_op(
     name='keys',

From a0bf30565dbe8b2fbdef9f8bdec3db7fb68f9db7 Mon Sep 17 00:00:00 2001
From: 97littleleaf11 <97littleleaf11@gmail.com>
Date: Wed, 7 Apr 2021 01:03:48 +0800
Subject: [PATCH 2/6] Add setdefault primitive

---
 mypyc/primitives/dict_ops.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/mypyc/primitives/dict_ops.py b/mypyc/primitives/dict_ops.py
index 62a4151c97f6..81577d1bcbc3 100644
--- a/mypyc/primitives/dict_ops.py
+++ b/mypyc/primitives/dict_ops.py
@@ -118,6 +118,14 @@
     c_function_name='CPyDict_GetWithNone',
     error_kind=ERR_MAGIC)
 
+# dict.setdefault(key, default)
+method_op(
+    name='setdefault',
+    arg_types=[dict_rprimitive, object_rprimitive, object_rprimitive],
+    return_type=object_rprimitive,
+    c_function_name='PyDict_SetDefault',
+    error_kind=ERR_MAGIC)
+
 # dict.keys()
 method_op(
     name='keys',

From 18fab0187afb30c9f849f4df6d43af8e18152458 Mon Sep 17 00:00:00 2001
From: 97littleleaf11 <97littleleaf11@gmail.com>
Date: Wed, 7 Apr 2021 01:18:13 +0800
Subject: [PATCH 3/6] Add test

---
 mypyc/test-data/fixtures/ir.py    |  1 +
 mypyc/test-data/irbuild-dict.test | 15 +++++++++++++++
 mypyc/test-data/run-dicts.test    |  7 +++++++
 3 files changed, 23 insertions(+)

diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py
index 70ecd9d9ab2e..fe4343c2d7fd 100644
--- a/mypyc/test-data/fixtures/ir.py
+++ b/mypyc/test-data/fixtures/ir.py
@@ -176,6 +176,7 @@ def values(self) -> Iterable[V]: pass
     def items(self) -> Iterable[Tuple[K, V]]: pass
     def clear(self) -> None: pass
     def copy(self) -> Dict[K, V]: pass
+    def setdefault(self, k: K, v: V) -> V: pass
 
 class set(Generic[T]):
     def __init__(self, i: Optional[Iterable[T]] = None) -> None: pass
diff --git a/mypyc/test-data/irbuild-dict.test b/mypyc/test-data/irbuild-dict.test
index 1225750b74df..214c123095f9 100644
--- a/mypyc/test-data/irbuild-dict.test
+++ b/mypyc/test-data/irbuild-dict.test
@@ -335,3 +335,18 @@ def f(d):
 L0:
     r0 = CPyDict_Copy(d)
     return r0
+
+[case testDictSetdefault]
+from typing import Dict
+def f(d: Dict[object, object]) -> object:
+    return d.setdefault('a', 'b')
+[out]
+def f(d):
+    d :: dict
+    r0, r1 :: str
+    r2 :: object
+L0:
+    r0 = 'a'
+    r1 = 'b'
+    r2 = PyDict_SetDefault(d, r0, r1)
+    return r2
diff --git a/mypyc/test-data/run-dicts.test b/mypyc/test-data/run-dicts.test
index 329d186874f4..fd0d32fb2967 100644
--- a/mypyc/test-data/run-dicts.test
+++ b/mypyc/test-data/run-dicts.test
@@ -217,3 +217,10 @@ def test_dict_copy() -> None:
     dd['a'] = 1
     assert dd.copy() == dd
     assert isinstance(dd.copy(), defaultdict)
+
+def test_dict_setdefault() -> None:
+    d: Dict[str, int] = {'a': 1, 'b': 2}
+    assert d.setdefault('a', 2) == 1
+    assert d.setdefault('c', 3) == 3
+    assert d['a'] == 1
+    assert d['c'] == 3

From 44836740b8b6758036b9aef9af51c48ac8ff8e2e Mon Sep 17 00:00:00 2001
From: 97littleleaf11 <97littleleaf11@gmail.com>
Date: Wed, 7 Apr 2021 21:31:08 +0800
Subject: [PATCH 4/6] Fix

---
 mypyc/lib-rt/CPy.h             |  1 +
 mypyc/lib-rt/dict_ops.c        |  4 ++++
 mypyc/primitives/dict_ops.py   | 10 ++++++++++
 mypyc/test-data/fixtures/ir.py |  2 +-
 mypyc/test-data/run-dicts.test |  3 +++
 5 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h
index 84b8bc81edec..821e5f704942 100644
--- a/mypyc/lib-rt/CPy.h
+++ b/mypyc/lib-rt/CPy.h
@@ -339,6 +339,7 @@ PyObject *CPyDict_GetItem(PyObject *dict, PyObject *key);
 int CPyDict_SetItem(PyObject *dict, PyObject *key, PyObject *value);
 PyObject *CPyDict_Get(PyObject *dict, PyObject *key, PyObject *fallback);
 PyObject *CPyDict_GetWithNone(PyObject *dict, PyObject *key);
+PyObject *CPyDict_SetDefaultWithNone(PyObject *dict, PyObject *key);
 PyObject *CPyDict_Build(Py_ssize_t size, ...);
 int CPyDict_Update(PyObject *dict, PyObject *stuff);
 int CPyDict_UpdateInDisplay(PyObject *dict, PyObject *stuff);
diff --git a/mypyc/lib-rt/dict_ops.c b/mypyc/lib-rt/dict_ops.c
index 4f831d0ba850..9c6359a1d802 100644
--- a/mypyc/lib-rt/dict_ops.c
+++ b/mypyc/lib-rt/dict_ops.c
@@ -67,6 +67,10 @@ PyObject *CPyDict_GetWithNone(PyObject *dict, PyObject *key) {
     return CPyDict_Get(dict, key, Py_None);
 }
 
+PyObject *CPyDict_SetDefaultWithNone(PyObject *dict, PyObject *key) {
+    return PyDict_SetDefault(dict, key, Py_None);
+}
+
 int CPyDict_SetItem(PyObject *dict, PyObject *key, PyObject *value) {
     if (PyDict_CheckExact(dict)) {
         return PyDict_SetItem(dict, key, value);
diff --git a/mypyc/primitives/dict_ops.py b/mypyc/primitives/dict_ops.py
index 81577d1bcbc3..8be7d0c2c9cb 100644
--- a/mypyc/primitives/dict_ops.py
+++ b/mypyc/primitives/dict_ops.py
@@ -124,6 +124,16 @@
     arg_types=[dict_rprimitive, object_rprimitive, object_rprimitive],
     return_type=object_rprimitive,
     c_function_name='PyDict_SetDefault',
+    is_borrowed=True,
+    error_kind=ERR_MAGIC)
+
+# dict.setdefault(key)
+method_op(
+    name='setdefault',
+    arg_types=[dict_rprimitive, object_rprimitive],
+    return_type=object_rprimitive,
+    c_function_name='CPyDict_SetDefaultWithNone',
+    is_borrowed=True,
     error_kind=ERR_MAGIC)
 
 # dict.keys()
diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py
index fe4343c2d7fd..e97f00dd71b6 100644
--- a/mypyc/test-data/fixtures/ir.py
+++ b/mypyc/test-data/fixtures/ir.py
@@ -176,7 +176,7 @@ def values(self) -> Iterable[V]: pass
     def items(self) -> Iterable[Tuple[K, V]]: pass
     def clear(self) -> None: pass
     def copy(self) -> Dict[K, V]: pass
-    def setdefault(self, k: K, v: V) -> V: pass
+    def setdefault(self, k: K, v: Optional[V] = None) -> Optional[V]: pass
 
 class set(Generic[T]):
     def __init__(self, i: Optional[Iterable[T]] = None) -> None: pass
diff --git a/mypyc/test-data/run-dicts.test b/mypyc/test-data/run-dicts.test
index fd0d32fb2967..6ec186b43d83 100644
--- a/mypyc/test-data/run-dicts.test
+++ b/mypyc/test-data/run-dicts.test
@@ -224,3 +224,6 @@ def test_dict_setdefault() -> None:
     assert d.setdefault('c', 3) == 3
     assert d['a'] == 1
     assert d['c'] == 3
+    assert d.setdefault('a') == 1
+    assert d.setdefault('e') == None
+    assert d.setdefault('e', 100) == None

From 74e0ba6009663b94f455fd0dbff9013699365fb8 Mon Sep 17 00:00:00 2001
From: 97littleleaf11 <97littleleaf11@gmail.com>
Date: Fri, 9 Apr 2021 07:54:38 +0800
Subject: [PATCH 5/6] Add test

---
 mypyc/lib-rt/CPy.h             |  1 +
 mypyc/lib-rt/dict_ops.c        |  9 ++++++++-
 mypyc/primitives/dict_ops.py   |  2 +-
 mypyc/test-data/run-dicts.test | 28 +++++++++++++++++++++++++++-
 4 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h
index 821e5f704942..953a496ad501 100644
--- a/mypyc/lib-rt/CPy.h
+++ b/mypyc/lib-rt/CPy.h
@@ -339,6 +339,7 @@ PyObject *CPyDict_GetItem(PyObject *dict, PyObject *key);
 int CPyDict_SetItem(PyObject *dict, PyObject *key, PyObject *value);
 PyObject *CPyDict_Get(PyObject *dict, PyObject *key, PyObject *fallback);
 PyObject *CPyDict_GetWithNone(PyObject *dict, PyObject *key);
+PyObject *CPyDict_SetDefault(PyObject *dict, PyObject *key, PyObject *value);
 PyObject *CPyDict_SetDefaultWithNone(PyObject *dict, PyObject *key);
 PyObject *CPyDict_Build(Py_ssize_t size, ...);
 int CPyDict_Update(PyObject *dict, PyObject *stuff);
diff --git a/mypyc/lib-rt/dict_ops.c b/mypyc/lib-rt/dict_ops.c
index 9c6359a1d802..c4dbe8d31c32 100644
--- a/mypyc/lib-rt/dict_ops.c
+++ b/mypyc/lib-rt/dict_ops.c
@@ -67,8 +67,15 @@ PyObject *CPyDict_GetWithNone(PyObject *dict, PyObject *key) {
     return CPyDict_Get(dict, key, Py_None);
 }
 
+PyObject *CPyDict_SetDefault(PyObject *dict, PyObject *key, PyObject *value) {
+    if (PyDict_CheckExact(dict)){
+        return PyDict_SetDefault(dict, key, value);
+    }
+    return PyObject_CallMethod(dict, "setdefault", "(OO)", key, value);
+}
+
 PyObject *CPyDict_SetDefaultWithNone(PyObject *dict, PyObject *key) {
-    return PyDict_SetDefault(dict, key, Py_None);
+    return CPyDict_SetDefault(dict, key, Py_None);
 }
 
 int CPyDict_SetItem(PyObject *dict, PyObject *key, PyObject *value) {
diff --git a/mypyc/primitives/dict_ops.py b/mypyc/primitives/dict_ops.py
index 8be7d0c2c9cb..267f9d79179d 100644
--- a/mypyc/primitives/dict_ops.py
+++ b/mypyc/primitives/dict_ops.py
@@ -123,7 +123,7 @@
     name='setdefault',
     arg_types=[dict_rprimitive, object_rprimitive, object_rprimitive],
     return_type=object_rprimitive,
-    c_function_name='PyDict_SetDefault',
+    c_function_name='CPyDict_SetDefault',
     is_borrowed=True,
     error_kind=ERR_MAGIC)
 
diff --git a/mypyc/test-data/run-dicts.test b/mypyc/test-data/run-dicts.test
index 6ec186b43d83..89188e09b1f1 100644
--- a/mypyc/test-data/run-dicts.test
+++ b/mypyc/test-data/run-dicts.test
@@ -196,7 +196,7 @@ else:
 
 [case testDictMethods]
 from collections import defaultdict
-from typing import Dict
+from typing import Dict, Optional
 
 def test_dict_clear() -> None:
     d = {'a': 1, 'b': 2}
@@ -218,12 +218,38 @@ def test_dict_copy() -> None:
     assert dd.copy() == dd
     assert isinstance(dd.copy(), defaultdict)
 
+class MyDict(dict):
+    def __init__(self, *args, **kwargs):
+        self.update(*args, **kwargs)
+
+    def setdefault(self, k, v=None):
+        if v is None:
+            if k in self.keys():
+                return self[k]
+            else:
+                return None
+        else:
+            return super().setdefault(k, v) + 10
+
 def test_dict_setdefault() -> None:
     d: Dict[str, int] = {'a': 1, 'b': 2}
     assert d.setdefault('a', 2) == 1
+    assert d.setdefault('b', 2) == 2
     assert d.setdefault('c', 3) == 3
     assert d['a'] == 1
     assert d['c'] == 3
     assert d.setdefault('a') == 1
     assert d.setdefault('e') == None
     assert d.setdefault('e', 100) == None
+
+def test_dict_subclass_setdefault() -> None:
+    d = MyDict()
+    d['a'] = 1
+    assert d.setdefault('a', 2) == 11
+    assert d.setdefault('b', 2) == 12
+    assert d.setdefault('c', 3) == 13
+    assert d['a'] == 1
+    assert d['c'] == 3
+    assert d.setdefault('a') == 1
+    assert d.setdefault('e') == None
+    assert d.setdefault('e', 100) == 110

From 8fa225f2fc5528c233b41a456d880e54d29266aa Mon Sep 17 00:00:00 2001
From: 97littleleaf11 <97littleleaf11@gmail.com>
Date: Fri, 9 Apr 2021 09:40:47 +0800
Subject: [PATCH 6/6] Fix IR test

---
 mypyc/test-data/irbuild-dict.test | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mypyc/test-data/irbuild-dict.test b/mypyc/test-data/irbuild-dict.test
index 214c123095f9..da08ed79d5bd 100644
--- a/mypyc/test-data/irbuild-dict.test
+++ b/mypyc/test-data/irbuild-dict.test
@@ -348,5 +348,5 @@ def f(d):
 L0:
     r0 = 'a'
     r1 = 'b'
-    r2 = PyDict_SetDefault(d, r0, r1)
+    r2 = CPyDict_SetDefault(d, r0, r1)
     return r2