Add writing functionality for dataframes

rdata/_write.py

@@ -1,17 +1,21 @@
 """Functions to perform conversion and unparsing in one step."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from .conversion import build_r_data, convert_to_r_object, convert_to_r_object_for_rda
-from .conversion.to_r import DEFAULT_FORMAT_VERSION
+from .conversion import (
+    DEFAULT_CONSTRUCTOR_DICT,
+    DEFAULT_FORMAT_VERSION,
+    convert_python_to_r_data,
+)
 from .unparser import unparse_file
 
 if TYPE_CHECKING:
     import os
     from typing import Any
 
-    from .conversion.to_r import Encoding
+    from .conversion.to_r import ConstructorDict, Encoding
     from .unparser import Compression, FileFormat
 
 
@@ -23,14 +27,12 @@ def write_rds(
     compression: Compression = "gzip",
     encoding: Encoding = "utf-8",
     format_version: int = DEFAULT_FORMAT_VERSION,
+    constructor_dict: ConstructorDict = DEFAULT_CONSTRUCTOR_DICT,
 ) -> None:
     """
     Write an RDS file.
 
-    This is a convenience function that wraps
-    :func:`rdata.conversion.convert_to_r_object`,
-    :func:`rdata.conversion.build_r_data`,
-    and :func:`rdata.unparser.unparse_file`,
+    This is a convenience function that wraps conversion and unparsing
     as it is the common use case.
 
     Args:
@@ -40,6 +42,7 @@ def write_rds(
         compression: Compression.
         encoding: Encoding to be used for strings within data.
         format_version: File format version.
+        constructor_dict: Dictionary mapping Python types to R classes.
 
     See Also:
         :func:`write_rda`: Similar function that writes an RDA or RDATA file.
@@ -52,15 +55,13 @@ def write_rds(
         >>> data = ["hello", 1, 2.2, 3.3+4.4j]
         >>> rdata.write_rds("test.rds", data)
     """
-    r_object = convert_to_r_object(
+    r_data = convert_python_to_r_data(
         data,
         encoding=encoding,
-    )
-    r_data = build_r_data(
-        r_object,
-        encoding=encoding,
         format_version=format_version,
+        constructor_dict=constructor_dict,
     )
+
     unparse_file(
         path,
         r_data,
@@ -78,14 +79,12 @@ def write_rda(
     compression: Compression = "gzip",
     encoding: Encoding = "utf-8",
     format_version: int = DEFAULT_FORMAT_VERSION,
+    constructor_dict: ConstructorDict = DEFAULT_CONSTRUCTOR_DICT,
 ) -> None:
     """
     Write an RDA or RDATA file.
 
-    This is a convenience function that wraps
-    :func:`rdata.conversion.convert_to_r_object_for_rda`,
-    :func:`rdata.conversion.build_r_data`,
-    and :func:`rdata.unparser.unparse_file`,
+    This is a convenience function that wraps conversion and unparsing
     as it is the common use case.
 
     Args:
@@ -95,6 +94,7 @@ def write_rda(
         compression: Compression.
         encoding: Encoding to be used for strings within data.
         format_version: File format version.
+        constructor_dict: Dictionary mapping Python types to R classes.
 
     See Also:
         :func:`write_rds`: Similar function that writes an RDS file.
@@ -107,15 +107,14 @@ def write_rda(
         >>> data = {"name": "hello", "values": [1, 2.2, 3.3+4.4j]}
         >>> rdata.write_rda("test.rda", data)
     """
-    r_object = convert_to_r_object_for_rda(
+    r_data = convert_python_to_r_data(
         data,
         encoding=encoding,
-    )
-    r_data = build_r_data(
-        r_object,
-        encoding=encoding,
         format_version=format_version,
+        constructor_dict=constructor_dict,
+        file_type="rda",
     )
+
     unparse_file(
         path,
         r_data,

rdata/conversion/__init__.py

@@ -1,4 +1,5 @@
 """Utilities for converting R objects to Python ones."""
+
 from ._conversion import (
     DEFAULT_CLASS_MAP as DEFAULT_CLASS_MAP,
     Converter as Converter,
@@ -25,7 +26,9 @@
     ts_constructor as ts_constructor,
 )
 from .to_r import (
-    build_r_data as build_r_data,
-    convert_to_r_object as convert_to_r_object,
-    convert_to_r_object_for_rda as convert_to_r_object_for_rda,
+    DEFAULT_CONSTRUCTOR_DICT as DEFAULT_CONSTRUCTOR_DICT,
+    DEFAULT_FORMAT_VERSION as DEFAULT_FORMAT_VERSION,
+    ConverterFromPythonToR as ConverterFromPythonToR,
+    convert_python_to_r_data as convert_python_to_r_data,
+    convert_python_to_r_object as convert_python_to_r_object,
 )

rdata/conversion/_conversion.py

@@ -394,20 +394,38 @@ def convert_array(
     return value  # type: ignore [no-any-return]
 
 
-R_INT_MIN = -2**31
-
-
 def _dataframe_column_transform(source: Any) -> Any:  # noqa: ANN401
 
     if isinstance(source, np.ndarray):
+        dtype: Any
         if np.issubdtype(source.dtype, np.integer):
-            return pd.Series(source, dtype=pd.Int32Dtype()).array
-
-        if np.issubdtype(source.dtype, np.bool_):
-            return pd.Series(source, dtype=pd.BooleanDtype()).array
+            dtype = pd.Int32Dtype()
+        elif np.issubdtype(source.dtype, np.floating):
+            # We return the numpy array here, which keeps
+            # R_FLOAT_NA, np.nan, and other NaNs as they were originally in the file.
+            # Users can then decide if they prefer to interpret
+            # only R_FLOAT_NA or all NaNs as "missing".
+            return source
+            # This would create an array with all NaNs as "missing":
+            # dtype = pd.Float64Dtype()  # noqa: ERA001
+            # This would create an array with only R_FLOAT_NA as "missing":
+            # from rdata.missing import is_na  # noqa: ERA001
+            # return pd.arrays.FloatingArray(source, is_na(source))  # noqa: ERA001
+        elif np.issubdtype(source.dtype, np.complexfloating):
+            # There seems to be no pandas type for complex array
+            return source
+        elif np.issubdtype(source.dtype, np.bool_):
+            dtype = pd.BooleanDtype()
+        elif np.issubdtype(source.dtype, np.str_):
+            dtype = pd.StringDtype()
+        elif np.issubdtype(source.dtype, np.object_):
+            for value in source:
+                assert isinstance(value, str) or value is None
+            dtype = pd.StringDtype()
+        else:
+            return source
 
-        if np.issubdtype(source.dtype, np.str_):
-            return pd.Series(source, dtype=pd.StringDtype()).array
+        return pd.Series(source, dtype=dtype).array
 
     return source
 
@@ -430,7 +448,7 @@ def dataframe_constructor(
             and isinstance(row_names, np.ma.MaskedArray)
             and row_names.mask[0]
         )
-        else tuple(row_names)
+        else row_names
     )
 
     return pd.DataFrame(obj, columns=obj, index=index)