fix: make DataFrame __getattr__ and __setattr__ more robust to subclassing (#1352)

* fix: make `DataFrame` `__getattr__` and `__setattr__` more robust to subclassing

* use _block as an initialized indicator

Author: tswast

bigframes/dataframe.py

@@ -118,6 +118,8 @@ def __init__(
     ):
         global bigframes
 
+        self._query_job: Optional[bigquery.QueryJob] = None
+
         if copy is not None and not copy:
             raise ValueError(
                 f"DataFrame constructor only supports copy=True. {constants.FEEDBACK_LINK}"
@@ -182,7 +184,6 @@ def __init__(
             if dtype:
                 bf_dtype = bigframes.dtypes.bigframes_type(dtype)
                 block = block.multi_apply_unary_op(ops.AsTypeOp(to_type=bf_dtype))
-            self._block = block
 
         else:
             import bigframes.pandas
@@ -194,10 +195,14 @@ def __init__(
                 dtype=dtype,  # type:ignore
             )
             if session:
-                self._block = session.read_pandas(pd_dataframe)._get_block()
+                block = session.read_pandas(pd_dataframe)._get_block()
             else:
-                self._block = bigframes.pandas.read_pandas(pd_dataframe)._get_block()
-        self._query_job: Optional[bigquery.QueryJob] = None
+                block = bigframes.pandas.read_pandas(pd_dataframe)._get_block()
+
+        # We use _block as an indicator in __getattr__ and __setattr__ to see
+        # if the object is fully initialized, so make sure we set the _block
+        # attribute last.
+        self._block = block
         self._block.session._register_object(self)
 
     def __dir__(self):
@@ -625,13 +630,17 @@ def _getitem_bool_series(self, key: bigframes.series.Series) -> DataFrame:
         return DataFrame(block)
 
     def __getattr__(self, key: str):
-        # Protect against recursion errors with uninitialized DataFrame
-        # objects. See:
+        # To allow subclasses to set private attributes before the class is
+        # fully initialized, protect against recursion errors with
+        # uninitialized DataFrame objects. Note: this comes at the downside
+        # that columns with a leading `_` won't be treated as columns.
+        #
+        # See:
         # https://github.com/googleapis/python-bigquery-dataframes/issues/728
         # and
         # https://nedbatchelder.com/blog/201010/surprising_getattr_recursion.html
         if key == "_block":
-            raise AttributeError("_block")
+            raise AttributeError(key)
 
         if key in self._block.column_labels:
             return self.__getitem__(key)
@@ -651,26 +660,36 @@ def __getattr__(self, key: str):
         raise AttributeError(key)
 
     def __setattr__(self, key: str, value):
-        if key in ["_block", "_query_job"]:
+        if key == "_block":
+            object.__setattr__(self, key, value)
+            return
+
+        # To allow subclasses to set private attributes before the class is
+        # fully initialized, assume anything set before `_block` is initialized
+        # is a regular attribute.
+        if not hasattr(self, "_block"):
             object.__setattr__(self, key, value)
             return
-        # Can this be removed???
+
+        # If someone has a column named the same as a normal attribute
+        # (e.g. index), we want to set the normal attribute, not the column.
+        # To do that, check if there is a normal attribute by using
+        # __getattribute__ (not __getattr__, because that includes columns).
+        # If that returns a value without raising, then we know this is a
+        # normal attribute and we should prefer that.
         try:
-            # boring attributes go through boring old path
             object.__getattribute__(self, key)
             return object.__setattr__(self, key, value)
         except AttributeError:
             pass
 
-        # if this fails, go on to more involved attribute setting
-        # (note that this matches __getattr__, above).
-        try:
-            if key in self.columns:
-                self[key] = value
-            else:
-                object.__setattr__(self, key, value)
-        # Can this be removed?
-        except (AttributeError, TypeError):
+        # If we made it here, then we know that it's not a regular attribute
+        # already, so it might be a column to update. Note: we don't allow
+        # adding new columns using __setattr__, only __setitem__, that way we
+        # can still add regular new attributes.
+        if key in self._block.column_labels:
+            self[key] = value
+        else:
             object.__setattr__(self, key, value)
 
     def __repr__(self) -> str:

tests/unit/test_dataframe.py

@@ -41,6 +41,15 @@ def test_dataframe_repr_with_uninitialized_object():
     assert "DataFrame" in got
 
 
+def test_dataframe_setattr_with_uninitialized_object():
+    """Ensures DataFrame can be subclassed without trying to set attributes as columns."""
+    # Avoid calling __init__ since it might be called later in a subclass.
+    # https://stackoverflow.com/a/6384982/101923
+    dataframe = bigframes.dataframe.DataFrame.__new__(bigframes.dataframe.DataFrame)
+    dataframe.lineage = "my-test-value"
+    assert dataframe.lineage == "my-test-value"  # Should just be a regular attribute.
+
+
 def test_dataframe_to_gbq_invalid_destination(monkeypatch: pytest.MonkeyPatch):
     dataframe = resources.create_dataframe(monkeypatch)