Parser v2

can_write_pdb(df)

Check whether atom data can be represented in PDB format without truncation.

For mmCIF-derived DataFrames, this enforces classic PDB limits on: atom serial numbers, chain identifiers and residue sequence numbers.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame with atom records and a format attribute set to "PDB" or "mmCIF".	required

Returns:

Name	Type	Description
bool	bool	True if the DataFrame fits PDB constraints, False otherwise.

Source code in src/rnapolis/parser_v2.py

def can_write_pdb(df: pd.DataFrame) -> bool:
    """Check whether atom data can be represented in PDB format without truncation.

    For mmCIF-derived DataFrames, this enforces classic PDB limits on:
    atom serial numbers, chain identifiers and residue sequence numbers.

    Args:
        df (pd.DataFrame): DataFrame with atom records and a ``format`` attribute
            set to ``"PDB"`` or ``"mmCIF"``.

    Returns:
        bool: True if the DataFrame fits PDB constraints, False otherwise.
    """
    format_type = df.attrs.get("format")

    if format_type == "PDB":
        # Assume data originally from PDB already fits PDB constraints
        return True

    if df.empty:
        # An empty DataFrame can be represented as an empty PDB file
        return True

    if format_type == "mmCIF":
        # Check serial number (id)
        # Convert to numeric first to handle potential categorical type and NaNs
        if "id" not in df.columns or (
            pd.to_numeric(df["id"], errors="coerce").max() > 99999
        ):
            return False

        # Check chain ID (auth_asym_id) length
        if "auth_asym_id" not in df.columns or (
            df["auth_asym_id"].dropna().astype(str).str.len().max() > 1
        ):
            return False

        # Check residue sequence number (auth_seq_id)
        if "auth_seq_id" not in df.columns or (
            pd.to_numeric(df["auth_seq_id"], errors="coerce").max() > 9999
        ):
            return False

        # All checks passed for mmCIF
        return True

    # If format is unknown or not PDB/mmCIF, assume it cannot be safely written
    return False

fit_to_pdb(df)

Try to renumber chains, residues and atoms so that data fits PDB limits.

If the DataFrame already satisfies PDB constraints, it is returned unchanged. Otherwise, the function checks feasibility (atoms, chains, residues/chain) and, if possible, remaps chain IDs, residue numbers and serials to safe ranges.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame with atom records and a valid format attribute ("PDB" or "mmCIF").	required

Returns:

Type	Description
DataFrame	pd.DataFrame: New DataFrame adjusted to comply with PDB constraints and
DataFrame	with format attribute set to "PDB".

Raises:

Type	Description
ValueError	If the data cannot be made to fit PDB limits (too many atoms, chains or residues per chain, or unsupported format).

Source code in src/rnapolis/parser_v2.py

def fit_to_pdb(df: pd.DataFrame) -> pd.DataFrame:
    """Try to renumber chains, residues and atoms so that data fits PDB limits.

    If the DataFrame already satisfies PDB constraints, it is returned unchanged.
    Otherwise, the function checks feasibility (atoms, chains, residues/chain)
    and, if possible, remaps chain IDs, residue numbers and serials to safe
    ranges.

    Args:
        df (pd.DataFrame): DataFrame with atom records and a valid ``format``
            attribute (``\"PDB\"`` or ``\"mmCIF\"``).

    Returns:
        pd.DataFrame: New DataFrame adjusted to comply with PDB constraints and
        with ``format`` attribute set to ``"PDB"``.

    Raises:
        ValueError: If the data cannot be made to fit PDB limits (too many
            atoms, chains or residues per chain, or unsupported format).
    """
    format_type = df.attrs.get("format")

    if not format_type:
        raise ValueError("DataFrame format attribute is not set.")

    if can_write_pdb(df):
        return df

    # Determine column names based on format
    if format_type == "PDB":
        serial_col = "serial"
        chain_col = "chainID"
        resseq_col = "resSeq"
        icode_col = "iCode"
    elif format_type == "mmCIF":
        serial_col = "id"
        chain_col = "auth_asym_id"
        resseq_col = "auth_seq_id"
        icode_col = "pdbx_PDB_ins_code"
    else:
        raise ValueError(f"Unsupported DataFrame format: {format_type}")

    # --- Feasibility Checks ---
    if chain_col not in df.columns:
        raise ValueError(f"Missing required chain column: {chain_col}")
    if resseq_col not in df.columns:
        raise ValueError(f"Missing required residue sequence column: {resseq_col}")

    unique_chains = df[chain_col].unique()
    num_chains = len(unique_chains)
    total_atoms = len(df)
    max_pdb_serial = 99999
    max_pdb_residue = 9999
    available_chain_ids = list(
        string.ascii_uppercase + string.ascii_lowercase + string.digits
    )
    max_pdb_chains = len(available_chain_ids)

    # Check 1: Total atoms + TER lines <= 99999
    if total_atoms + num_chains > max_pdb_serial:
        raise ValueError(
            f"Cannot fit to PDB: Total atoms ({total_atoms}) + TER lines ({num_chains}) exceeds PDB limit ({max_pdb_serial})."
        )

    # Check 2: Number of chains <= 62
    if num_chains > max_pdb_chains:
        raise ValueError(
            f"Cannot fit to PDB: Number of unique chains ({num_chains}) exceeds PDB limit ({max_pdb_chains})."
        )

    # Check 3: Max residues per chain <= 9999
    # More accurate check: group by chain, then count unique (resSeq, iCode) tuples
    # Use a temporary structure to avoid modifying the original df
    check_df = pd.DataFrame(
        {
            "chain": df[chain_col],
            "resSeq": df[resseq_col],
            "iCode": df[icode_col].fillna("") if icode_col in df.columns else "",
        }
    )
    residue_counts = check_df.groupby("chain").apply(
        lambda x: x[["resSeq", "iCode"]].drop_duplicates().shape[0]
    )
    max_residues_per_chain = residue_counts.max() if not residue_counts.empty else 0

    if max_residues_per_chain > max_pdb_residue:
        raise ValueError(
            f"Cannot fit to PDB: Maximum residues in a single chain ({max_residues_per_chain}) exceeds PDB limit ({max_pdb_residue})."
        )

    # --- Perform Fitting ---
    df_fitted = df.copy()

    # 1. Rename Chains
    chain_mapping = {
        orig_chain: available_chain_ids[i] for i, orig_chain in enumerate(unique_chains)
    }
    df_fitted[chain_col] = df_fitted[chain_col].map(chain_mapping)
    # Ensure the chain column is treated as string/object after mapping
    df_fitted[chain_col] = df_fitted[chain_col].astype(object)

    # 2. Renumber Residues within each new chain
    new_resseq_col = "new_resSeq"  # Temporary column for new numbering
    df_fitted[new_resseq_col] = -1  # Initialize

    all_new_res_maps = {}
    for new_chain_id, group in df_fitted.groupby(chain_col):
        # Identify unique original residues (seq + icode) in order of appearance
        original_residues = group[[resseq_col, icode_col]].drop_duplicates()
        # Create mapping: (orig_resSeq, orig_iCode) -> new_resSeq (1-based)
        residue_mapping = {
            tuple(res): i + 1
            for i, res in enumerate(original_residues.itertuples(index=False))
        }
        all_new_res_maps[new_chain_id] = residue_mapping

        # Apply mapping to the group
        res_indices = group.set_index([resseq_col, icode_col]).index
        df_fitted.loc[group.index, new_resseq_col] = res_indices.map(residue_mapping)

    # Replace original residue number and clear insertion code
    df_fitted[resseq_col] = df_fitted[new_resseq_col]
    df_fitted[icode_col] = None  # Insertion codes are now redundant
    df_fitted.drop(columns=[new_resseq_col], inplace=True)
    # Convert resseq_col back to Int64 if it was before, handling potential NaNs if any step failed
    df_fitted[resseq_col] = df_fitted[resseq_col].astype("Int64")

    # 3. Renumber Atom Serials
    new_serial_col = "new_serial"
    df_fitted[new_serial_col] = -1  # Initialize
    current_serial = 0
    last_chain_id_for_serial = None

    # Iterate in the potentially re-sorted order after grouping/mapping
    # Ensure stable sort order for consistent serial numbering
    df_fitted.sort_index(
        inplace=True
    )  # Sort by original index to maintain original atom order as much as possible

    for index, row in df_fitted.iterrows():
        current_chain_id = row[chain_col]
        if (
            last_chain_id_for_serial is not None
            and current_chain_id != last_chain_id_for_serial
        ):
            current_serial += 1  # Increment for TER line

        current_serial += 1
        if current_serial > max_pdb_serial:
            # This should have been caught by the initial check, but is a safeguard
            raise ValueError("Serial number exceeded PDB limit during renumbering.")

        df_fitted.loc[index, new_serial_col] = current_serial
        last_chain_id_for_serial = current_chain_id

    # Replace original serial number
    df_fitted[serial_col] = df_fitted[new_serial_col]
    df_fitted.drop(columns=[new_serial_col], inplace=True)
    # Convert serial_col back to Int64
    df_fitted[serial_col] = df_fitted[serial_col].astype("Int64")

    # Update attributes and column types for PDB compatibility
    df_fitted.attrs["format"] = "PDB"

    # Ensure final column types match expected PDB output (especially categories)
    # Reapply categorical conversion as some operations might change dtypes
    pdb_categorical_cols = [
        "record_type",
        "name",
        "altLoc",
        "resName",
        chain_col,
        "element",
        "charge",
        icode_col,
    ]
    if "record_type" not in df_fitted.columns and "group_PDB" in df_fitted.columns:
        df_fitted.rename(
            columns={"group_PDB": "record_type"}, inplace=True
        )  # Ensure correct name

    for col in pdb_categorical_cols:
        if col in df_fitted.columns:
            # Handle None explicitly before converting to category if needed
            if df_fitted[col].isnull().any():
                df_fitted[col] = (
                    df_fitted[col].astype(object).fillna("")
                )  # Fill None with empty string for category
            df_fitted[col] = df_fitted[col].astype("category")

    # Rename columns if necessary from mmCIF to PDB standard names
    rename_map = {
        "id": "serial",
        "auth_asym_id": "chainID",
        "auth_seq_id": "resSeq",
        "pdbx_PDB_ins_code": "iCode",
        "label_atom_id": "name",  # Prefer label_atom_id if auth_atom_id not present? PDB uses 'name'
        "label_comp_id": "resName",  # Prefer label_comp_id if auth_comp_id not present? PDB uses 'resName'
        "type_symbol": "element",
        "pdbx_formal_charge": "charge",
        "Cartn_x": "x",
        "Cartn_y": "y",
        "Cartn_z": "z",
        "B_iso_or_equiv": "tempFactor",
        "group_PDB": "record_type",
        "pdbx_PDB_model_num": "model",
        # Add mappings for auth_atom_id -> name, auth_comp_id -> resName if needed,
        # deciding on precedence if both label_* and auth_* exist.
        # Current write_pdb prioritizes auth_* when reading mmCIF, so map those.
        "auth_atom_id": "name",
        "auth_comp_id": "resName",
    }

    # Only rename columns that actually exist in the DataFrame
    actual_rename_map = {k: v for k, v in rename_map.items() if k in df_fitted.columns}
    df_fitted.rename(columns=actual_rename_map, inplace=True)

    # Ensure essential PDB columns exist, even if empty, if they were created during fitting
    pdb_essential_cols = [
        "record_type",
        "serial",
        "name",
        "altLoc",
        "resName",
        "chainID",
        "resSeq",
        "iCode",
        "x",
        "y",
        "z",
        "occupancy",
        "tempFactor",
        "element",
        "charge",
        "model",
    ]
    for col in pdb_essential_cols:
        if col not in df_fitted.columns:
            # This case might occur if input mmCIF was missing fundamental columns mapped to PDB essentials
            # Decide on default value or raise error. Adding empty series for now.
            df_fitted[col] = pd.Series(
                dtype="object"
            )  # Add as object to handle potential None/mixed types initially

    # Re-order columns to standard PDB order for clarity
    final_pdb_order = [col for col in pdb_essential_cols if col in df_fitted.columns]
    other_cols = [col for col in df_fitted.columns if col not in final_pdb_order]
    df_fitted = df_fitted[final_pdb_order + other_cols]

    # --- Final Type Conversions for PDB format ---
    # Convert numeric columns (similar to parse_pdb_atoms)
    pdb_numeric_columns = [
        "serial",
        "resSeq",
        "x",
        "y",
        "z",
        "occupancy",
        "tempFactor",
        "model",
    ]
    for col in pdb_numeric_columns:
        if col in df_fitted.columns:
            # Use Int64 for integer-like columns that might have been NaN during processing
            if col in ["serial", "resSeq", "model"]:
                df_fitted[col] = pd.to_numeric(df_fitted[col], errors="coerce").astype(
                    "Int64"
                )
            else:  # Floats
                df_fitted[col] = pd.to_numeric(df_fitted[col], errors="coerce")

    # Convert categorical columns (similar to parse_pdb_atoms)
    # Note: chainID and iCode were already handled during fitting/renaming
    pdb_categorical_columns_final = [
        "record_type",
        "name",
        "altLoc",
        "resName",
        "chainID",  # Already category, but ensure consistency
        "iCode",  # Already category, but ensure consistency
        "element",
        "charge",
    ]
    for col in pdb_categorical_columns_final:
        if col in df_fitted.columns:
            # Ensure the column is categorical first
            if not pd.api.types.is_categorical_dtype(df_fitted[col]):
                # Convert non-categorical columns, handling potential NaNs
                if df_fitted[col].isnull().any():
                    df_fitted[col] = (
                        df_fitted[col].astype(object).fillna("").astype("category")
                    )
                else:
                    df_fitted[col] = df_fitted[col].astype("category")
            else:
                # If already categorical, check if '' needs to be added before fillna
                has_nans = df_fitted[col].isnull().any()
                if has_nans and "" not in df_fitted[col].cat.categories:
                    # Add '' category explicitly
                    df_fitted[col] = df_fitted[col].cat.add_categories([""])

                # Fill None/NaN with empty string (now safe)
                if has_nans:
                    df_fitted[col].fillna("", inplace=True)

    return df_fitted

parse_cif_atoms(content)

Parse mmCIF content and extract atom_site records into a DataFrame.

Supports string input, StringIO and file-like objects with a name attribute. Missing values marked as ? or . are converted to None and selected columns are cast to numeric or categorical dtypes.

Parameters:

Name	Type	Description	Default
content	Union[str, IO[str]]	mmCIF content as a string, StringIO or an open file-like object.	required

Returns:

Type	Description
DataFrame	pd.DataFrame: DataFrame with parsed atom_site records and a format attribute set to "mmCIF". Empty if no atom_site category is present.

Source code in src/rnapolis/parser_v2.py

def parse_cif_atoms(content: Union[str, IO[str]]) -> pd.DataFrame:
    """Parse mmCIF content and extract ``atom_site`` records into a DataFrame.

    Supports string input, StringIO and file-like objects with a ``name``
    attribute. Missing values marked as ``?`` or ``.`` are converted to None
    and selected columns are cast to numeric or categorical dtypes.

    Args:
        content (Union[str, IO[str]]): mmCIF content as a string, StringIO or
            an open file-like object.

    Returns:
        pd.DataFrame: DataFrame with parsed ``atom_site`` records and a
            ``format`` attribute set to ``"mmCIF"``. Empty if no ``atom_site``
            category is present.
    """
    adapter = IoAdapterPy()

    if isinstance(content, str):
        cif_text = content
    elif isinstance(content, io.StringIO):
        content.seek(0)
        cif_text = content.read()
    elif hasattr(content, "read"):
        content.seek(0)
        cif_text = content.read()
        if isinstance(cif_text, bytes):
            cif_text = cif_text.decode("utf-8")
    else:
        raise TypeError(
            "Unsupported input type for parse_cif_atoms. Expected str, file-like object, or StringIO."
        )

    cif_text = re.sub(r"(^\s*data_)\s*$", r"\1unnamed", cif_text, flags=re.MULTILINE)
    for axis in ("x", "y", "z"):
        cif_text = re.sub(
            rf"(_atom_site\.)cartn_{axis}\b",
            rf"\1Cartn_{axis}",
            cif_text,
            flags=re.IGNORECASE,
        )

    with tempfile.NamedTemporaryFile(
        mode="w+", suffix=".cif", delete=False
    ) as temp_file:
        temp_file.write(cif_text)
        temp_file_path = temp_file.name
    try:
        data = adapter.readFile(temp_file_path)
    finally:
        os.remove(temp_file_path)

    # Get the atom_site category
    category = data[0].getObj("atom_site")

    if not category:
        # Return empty DataFrame if no atom_site category found
        return pd.DataFrame()

    # Extract attribute names and data rows
    attributes = category.getAttributeList()
    rows = category.getRowList()

    # Create a list of dictionaries for each atom
    records = []
    for row in rows:
        record = {}
        for attr, value in zip(attributes, row):
            # Store None if value indicates missing data ('?' or '.')
            if value in ["?", "."]:
                record[attr] = None
            else:
                record[attr] = value
        records.append(record)

    # Create DataFrame from records
    df = pd.DataFrame(records)

    # Define columns based on mmCIF specification for atom_site
    float_cols = [
        "aniso_B[1][1]",
        "aniso_B[1][1]_esd",
        "aniso_B[1][2]",
        "aniso_B[1][2]_esd",
        "aniso_B[1][3]",
        "aniso_B[1][3]_esd",
        "aniso_B[2][2]",
        "aniso_B[2][2]_esd",
        "aniso_B[2][3]",
        "aniso_B[2][3]_esd",
        "aniso_B[3][3]",
        "aniso_B[3][3]_esd",
        "aniso_ratio",
        "aniso_U[1][1]",
        "aniso_U[1][1]_esd",
        "aniso_U[1][2]",
        "aniso_U[1][2]_esd",
        "aniso_U[1][3]",
        "aniso_U[1][3]_esd",
        "aniso_U[2][2]",
        "aniso_U[2][2]_esd",
        "aniso_U[2][3]",
        "aniso_U[2][3]_esd",
        "aniso_U[3][3]",
        "aniso_U[3][3]_esd",
        "B_equiv_geom_mean",
        "B_equiv_geom_mean_esd",
        "B_iso_or_equiv",
        "B_iso_or_equiv_esd",
        "Cartn_x",
        "Cartn_x_esd",
        "Cartn_y",
        "Cartn_y_esd",
        "Cartn_z",
        "Cartn_z_esd",
        "fract_x",
        "fract_x_esd",
        "fract_y",
        "fract_y_esd",
        "fract_z",
        "fract_z_esd",
        "occupancy",
        "occupancy_esd",
        "U_equiv_geom_mean",
        "U_equiv_geom_mean_esd",
        "U_iso_or_equiv",
        "U_iso_or_equiv_esd",
    ]
    int_cols = [
        "attached_hydrogens",
        "label_seq_id",
        "symmetry_multiplicity",
        "pdbx_PDB_model_num",
        "pdbx_formal_charge",
        "pdbx_label_index",
    ]
    category_cols = [
        "auth_asym_id",
        "auth_atom_id",
        "auth_comp_id",
        "auth_seq_id",
        "calc_attached_atom",
        "calc_flag",
        "disorder_assembly",
        "disorder_group",
        "group_PDB",
        "id",
        "label_alt_id",
        "label_asym_id",
        "label_atom_id",
        "label_comp_id",
        "label_entity_id",
        "thermal_displace_type",
        "type_symbol",
        "pdbx_atom_ambiguity",
        "adp_type",
        "refinement_flags",
        "refinement_flags_adp",
        "refinement_flags_occupancy",
        "refinement_flags_posn",
        "pdbx_auth_alt_id",
        "pdbx_PDB_ins_code",
        "pdbx_PDB_residue_no",
        "pdbx_PDB_residue_name",
        "pdbx_PDB_strand_id",
        "pdbx_PDB_atom_name",
        "pdbx_auth_atom_name",
        "pdbx_auth_comp_id",
        "pdbx_auth_asym_id",
        "pdbx_auth_seq_id",
        "pdbx_tls_group_id",
        "pdbx_ncs_dom_id",
        "pdbx_group_NDB",
        "pdbx_atom_group",
        "pdbx_label_seq_num",
        "pdbx_not_in_asym",
        "pdbx_sifts_xref_db_name",
        "pdbx_sifts_xref_db_acc",
        "pdbx_sifts_xref_db_num",
        "pdbx_sifts_xref_db_res",
    ]

    # Convert columns to appropriate types
    for col in float_cols:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")

    for col in int_cols:
        if col in df.columns:
            # Use Int64 (nullable integer) to handle potential NaNs from coercion
            df[col] = pd.to_numeric(df[col], errors="coerce").astype("Int64")

    for col in category_cols:
        if col in df.columns:
            df[col] = df[col].astype("category")

    # Add format attribute to the DataFrame
    df.attrs["format"] = "mmCIF"

    return df

parse_cif_modres(content)

Parse mmCIF pdbx_struct_mod_residue category into a DataFrame.

Parameters:

Name	Type	Description	Default
content	Union[str, IO[str]]	mmCIF content as a string or file-like object.	required

Returns:

Type	Description
DataFrame	DataFrame with native mmCIF column names: auth_comp_id, auth_asym_id,
DataFrame	auth_seq_id, pdbx_PDB_ins_code, parent_comp_id, details.
DataFrame	Empty DataFrame with correct columns if no pdbx_struct_mod_residue found.

Source code in src/rnapolis/parser_v2.py

def parse_cif_modres(content: Union[str, IO[str]]) -> pd.DataFrame:
    """Parse mmCIF pdbx_struct_mod_residue category into a DataFrame.

    Args:
        content: mmCIF content as a string or file-like object.

    Returns:
        DataFrame with native mmCIF column names: auth_comp_id, auth_asym_id,
        auth_seq_id, pdbx_PDB_ins_code, parent_comp_id, details.
        Empty DataFrame with correct columns if no pdbx_struct_mod_residue found.
    """
    adapter = IoAdapterPy()

    if isinstance(content, str):
        cif_text = content
    elif isinstance(content, io.StringIO):
        content.seek(0)
        cif_text = content.read()
    elif hasattr(content, "read"):
        content.seek(0)
        cif_text = content.read()
        if isinstance(cif_text, bytes):
            cif_text = cif_text.decode("utf-8")
    else:
        raise TypeError("Unsupported input type for parse_cif_modres.")

    with tempfile.NamedTemporaryFile(
        mode="w+", suffix=".cif", delete=False
    ) as temp_file:
        temp_file.write(cif_text)
        temp_file_path = temp_file.name
    try:
        data = adapter.readFile(temp_file_path)
    finally:
        os.remove(temp_file_path)

    category = data[0].getObj("pdbx_struct_mod_residue") if data else None

    result_cols = [
        "auth_comp_id",
        "auth_asym_id",
        "auth_seq_id",
        "pdbx_PDB_ins_code",
        "parent_comp_id",
        "details",
    ]

    if not category:
        df = pd.DataFrame(columns=result_cols)
        df.attrs["format"] = "mmCIF"
        return df

    attributes = category.getAttributeList()
    rows = category.getRowList()

    records = []
    for row in rows:
        record = {}
        for attr, value in zip(attributes, row):
            record[attr] = None if value in ["?", "."] else value
        records.append(record)

    df = pd.DataFrame(records)

    for col in ["auth_comp_id", "auth_asym_id", "parent_comp_id"]:
        if col in df.columns:
            df[col] = df[col].astype("category")

    if "pdbx_PDB_ins_code" in df.columns:
        df["pdbx_PDB_ins_code"] = (
            df["pdbx_PDB_ins_code"].fillna("").astype(str).replace("None", "")
        )

    if "auth_seq_id" in df.columns:
        df["auth_seq_id"] = pd.to_numeric(df["auth_seq_id"], errors="coerce").astype(
            "Int64"
        )

    existing_cols = [col for col in result_cols if col in df.columns]
    df = df[existing_cols] if existing_cols else pd.DataFrame(columns=result_cols)

    df.attrs["format"] = "mmCIF"
    return df

parse_pdb_atoms(content)

Parse PDB content and extract ATOM/HETATM records into a DataFrame.

Handles both plain strings and file-like objects, collects atom-level fields and returns them as a typed pandas DataFrame.

Parameters:

Name	Type	Description	Default
content	Union[str, IO[str]]	PDB content as a string or an open text file-like object.	required

Returns:

Type	Description
DataFrame	pd.DataFrame: DataFrame with parsed ATOM/HETATM records and a format attribute set to "PDB".

Source code in src/rnapolis/parser_v2.py

def parse_pdb_atoms(content: Union[str, IO[str]]) -> pd.DataFrame:
    """Parse PDB content and extract ATOM/HETATM records into a DataFrame.

    Handles both plain strings and file-like objects, collects atom-level fields
    and returns them as a typed pandas DataFrame.

    Args:
        content (Union[str, IO[str]]): PDB content as a string or an open
            text file-like object.

    Returns:
        pd.DataFrame: DataFrame with parsed ATOM/HETATM records and a
            ``format`` attribute set to ``"PDB"``.
    """
    records = []

    # Handle both string content and file-like objects
    if isinstance(content, str):
        lines = content.splitlines()
    else:
        # Read all lines from the file-like object
        content.seek(0)  # Ensure we're at the beginning of the file
        lines = content.readlines()
        # Convert bytes to string if needed
        if isinstance(lines[0], bytes):
            lines = [line.decode("utf-8") for line in lines]

    current_model = 1
    for line in lines:
        record_type = line[:6].strip()

        # Check for MODEL record
        if record_type == "MODEL":
            try:
                current_model = int(line[10:14].strip())
            except ValueError:
                # Handle cases where MODEL record might be malformed
                pass  # Keep the previous model number
            continue

        # Only process ATOM and HETATM records
        if record_type not in ["ATOM", "HETATM"]:
            continue

        # Parse fields according to PDB format specification
        alt_loc = line[16:17].strip()
        icode = line[26:27].strip()
        element = line[76:78].strip()
        charge = line[78:80].strip()

        record = {
            "record_type": record_type,
            "serial": line[6:11].strip(),
            "name": line[12:16].strip(),
            "altLoc": None if not alt_loc else alt_loc,  # Store None if empty
            "resName": line[17:20].strip(),
            "chainID": line[21:22].strip(),
            "resSeq": line[22:26].strip(),
            "iCode": None if not icode else icode,  # Store None if empty
            "x": line[30:38].strip(),
            "y": line[38:46].strip(),
            "z": line[46:54].strip(),
            "occupancy": line[54:60].strip(),
            "tempFactor": line[60:66].strip(),
            "element": None if not element else element,  # Store None if empty
            "charge": None if not charge else charge,  # Store None if empty
            "model": current_model,  # Add the current model number
        }

        records.append(record)

    # Create DataFrame from records
    if not records:
        # Return empty DataFrame with correct columns if no records found
        return pd.DataFrame(
            columns=[
                "record_type",
                "serial",
                "name",
                "altLoc",
                "resName",
                "chainID",
                "resSeq",
                "iCode",
                "x",
                "y",
                "z",
                "occupancy",
                "tempFactor",
                "element",
                "charge",
                "model",
            ]
        )

    df = pd.DataFrame(records)

    # Convert numeric columns to appropriate types
    numeric_columns = [
        "serial",
        "resSeq",
        "x",
        "y",
        "z",
        "occupancy",
        "tempFactor",
        "model",
    ]
    for col in numeric_columns:
        df[col] = pd.to_numeric(df[col], errors="coerce")

    # Convert categorical columns
    categorical_columns = [
        "record_type",
        "name",
        "altLoc",
        "resName",
        "chainID",
        "element",
        "charge",
    ]
    for col in categorical_columns:
        df[col] = df[col].astype("category")

    # Add format attribute to the DataFrame
    df.attrs["format"] = "PDB"

    return df

parse_pdb_modres(content)

Parse PDB MODRES records into a DataFrame.

Parameters:

Name	Type	Description	Default
content	Union[str, IO[str]]	PDB content as a string or file-like object.	required

Returns:

Type	Description
DataFrame	DataFrame with columns: resName, chainID, seqNum, iCode, stdRes, comment.
DataFrame	Empty DataFrame with correct columns if no MODRES records found.

Source code in src/rnapolis/parser_v2.py

def parse_pdb_modres(content: Union[str, IO[str]]) -> pd.DataFrame:
    """Parse PDB MODRES records into a DataFrame.

    Args:
        content: PDB content as a string or file-like object.

    Returns:
        DataFrame with columns: resName, chainID, seqNum, iCode, stdRes, comment.
        Empty DataFrame with correct columns if no MODRES records found.
    """
    records = []

    if isinstance(content, str):
        lines = content.splitlines()
    else:
        content.seek(0)
        lines = content.readlines()
        if isinstance(lines[0], bytes):
            lines = [line.decode("utf-8") for line in lines]

    for line in lines:
        if not line.startswith("MODRES"):
            continue

        res_name = line[12:15].strip() if len(line) > 12 else ""
        chain_id = line[16:17].strip() if len(line) > 16 else ""
        seq_num = line[18:22].strip() if len(line) > 18 else ""
        i_code = line[22:23].strip() if len(line) > 22 else ""
        std_res = line[24:27].strip() if len(line) > 24 else ""
        comment = line[29:70].strip() if len(line) > 29 else ""

        record = {
            "resName": res_name,
            "chainID": chain_id,
            "seqNum": seq_num,
            "iCode": None if not i_code else i_code,
            "stdRes": std_res,
            "comment": None if not comment else comment,
        }
        records.append(record)

    if not records:
        df = pd.DataFrame(
            columns=["resName", "chainID", "seqNum", "iCode", "stdRes", "comment"]
        )
    else:
        df = pd.DataFrame(records)
        df["seqNum"] = pd.to_numeric(df["seqNum"], errors="coerce").astype("Int64")
        for col in ["resName", "chainID", "stdRes"]:
            df[col] = df[col].astype("category")
        df["iCode"] = df["iCode"].fillna("").astype(str).replace("nan", "")

    df.attrs["format"] = "PDB"
    return df

write_cif(df, output=None)

Write atom records stored in a DataFrame to mmCIF format.

Depending on the DataFrame format attribute, either passes through existing mmCIF-style columns or maps PDB-style columns to an atom_site category.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame with atom records, usually created by :func:parse_pdb_atoms or :func:parse_cif_atoms.	required
output	Union[str, TextIO, None]	File path, open text handle or None. If None, the mmCIF content is returned as a string.	None

Returns:

Type	Description
Union[str, None]	Union[str, None]: mmCIF content as a string if output is None,
Union[str, None]	otherwise None.

Source code in src/rnapolis/parser_v2.py

def write_cif(
    df: pd.DataFrame, output: Union[str, TextIO, None] = None
) -> Union[str, None]:
    """Write atom records stored in a DataFrame to mmCIF format.

    Depending on the DataFrame ``format`` attribute, either passes through
    existing mmCIF-style columns or maps PDB-style columns to an ``atom_site``
    category.

    Args:
        df (pd.DataFrame): DataFrame with atom records, usually created by
            :func:`parse_pdb_atoms` or :func:`parse_cif_atoms`.
        output (Union[str, TextIO, None], optional): File path, open text handle
            or None. If None, the mmCIF content is returned as a string.

    Returns:
        Union[str, None]: mmCIF content as a string if ``output`` is None,
        otherwise None.
    """
    # Get the format of the DataFrame
    format_type = df.attrs.get("format", "PDB")

    # Create a new DataContainer
    data_container = DataContainer("rnapolis")

    # Define the attributes for atom_site category
    if format_type == "mmCIF":
        # Use existing mmCIF attributes
        attributes = list(df.columns)
    else:  # PDB format
        # Map PDB columns to mmCIF attributes
        attributes = [
            "group_PDB",  # record_type
            "id",  # serial
            "type_symbol",  # element
            "label_atom_id",  # name
            "label_alt_id",  # altLoc
            "label_comp_id",  # resName
            "label_asym_id",  # chainID
            "label_entity_id",  # (generated)
            "label_seq_id",  # resSeq
            "pdbx_PDB_ins_code",  # iCode
            "Cartn_x",  # x
            "Cartn_y",  # y
            "Cartn_z",  # z
            "occupancy",  # occupancy
            "B_iso_or_equiv",  # tempFactor
            "pdbx_formal_charge",  # charge
            "auth_seq_id",  # resSeq
            "auth_comp_id",  # resName
            "auth_asym_id",  # chainID
            "auth_atom_id",  # name
            "pdbx_PDB_model_num",  # model
        ]

    # Prepare rows for the atom_site category
    rows = []

    for _, row in df.iterrows():
        if format_type == "mmCIF":
            # Use existing mmCIF data, converting None to '?' universally
            row_data = []
            for attr in attributes:
                value = row.get(attr)
                if pd.isna(value):
                    # Use '?' as the standard placeholder for missing values
                    row_data.append("?")
                else:
                    # Ensure all non-missing values are converted to string
                    row_data.append(str(value))
        else:  # PDB format
            # Map PDB data to mmCIF format, converting None to '.' or '?'
            entity_id = "1"  # Default entity ID
            model_num = str(int(row["model"]))

            # Pre-process optional fields for mmCIF placeholders
            element_val = "?" if pd.isna(row.get("element")) else str(row["element"])
            altloc_val = "." if pd.isna(row.get("altLoc")) else str(row["altLoc"])
            icode_val = "." if pd.isna(row.get("iCode")) else str(row["iCode"])
            charge_val = "." if pd.isna(row.get("charge")) else str(row["charge"])

            row_data = [
                str(row["record_type"]),  # group_PDB
                str(int(row["serial"])),  # id
                element_val,  # type_symbol
                str(row["name"]),  # label_atom_id
                altloc_val,  # label_alt_id
                str(row["resName"]),  # label_comp_id
                str(row["chainID"]),  # label_asym_id
                entity_id,  # label_entity_id
                str(int(row["resSeq"])),  # label_seq_id
                icode_val,  # pdbx_PDB_ins_code
                f"{float(row['x']):.3f}",  # Cartn_x
                f"{float(row['y']):.3f}",  # Cartn_y
                f"{float(row['z']):.3f}",  # Cartn_z
                f"{float(row['occupancy']):.2f}",  # occupancy
                f"{float(row['tempFactor']):.2f}",  # B_iso_or_equiv
                charge_val,  # pdbx_formal_charge
                str(int(row["resSeq"])),  # auth_seq_id
                str(row["resName"]),  # auth_comp_id
                str(row["chainID"]),  # auth_asym_id
                str(row["name"]),  # auth_atom_id
                model_num,  # pdbx_PDB_model_num
            ]

        rows.append(row_data)

    # Create the atom_site category
    atom_site_category = DataCategory("atom_site", attributes, rows)

    # Add the category to the data container
    data_container.append(atom_site_category)

    # Create an IoAdapter for writing
    adapter = IoAdapterPy()

    # Handle output
    if output is None:
        # Return as string - write to a temporary file and read it back
        with tempfile.NamedTemporaryFile(mode="w+", suffix=".cif") as temp_file:
            adapter.writeFile(temp_file.name, [data_container])
            temp_file.flush()
            temp_file.seek(0)
            return temp_file.read()
    elif isinstance(output, str):
        # Write to a file path
        adapter.writeFile(output, [data_container])
        return None
    else:
        # Write to a file-like object
        with tempfile.NamedTemporaryFile(mode="w+", suffix=".cif") as temp_file:
            adapter.writeFile(temp_file.name, [data_container])
            temp_file.flush()
            temp_file.seek(0)
            output.write(temp_file.read())
        return None

write_pdb(df, output=None)

Write atom records stored in a DataFrame to PDB format.

Supports both data originating from PDB and mmCIF (via the format attribute). Can either return a string with PDB content or write to a file.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame with atom records, typically produced by :func:parse_pdb_atoms or :func:parse_cif_atoms.	required
output	Union[str, TextIO, None]	File path, open text handle or None. If None, the function returns the PDB string.	None

Returns:

Type	Description
Union[str, None]	Union[str, None]: PDB content as a string if output is None,
Union[str, None]	otherwise None.

Source code in src/rnapolis/parser_v2.py

def write_pdb(
    df: pd.DataFrame, output: Union[str, TextIO, None] = None
) -> Union[str, None]:
    """Write atom records stored in a DataFrame to PDB format.

    Supports both data originating from PDB and mmCIF (via the ``format``
    attribute). Can either return a string with PDB content or write to a file.

    Args:
        df (pd.DataFrame): DataFrame with atom records, typically produced by
            :func:`parse_pdb_atoms` or :func:`parse_cif_atoms`.
        output (Union[str, TextIO, None], optional): File path, open text handle
            or None. If None, the function returns the PDB string.

    Returns:
        Union[str, None]: PDB content as a string if ``output`` is None,
        otherwise None.
    """
    buffer = io.StringIO()
    format_type = df.attrs.get("format", "PDB")  # Assume PDB if not specified

    last_model_num = None
    last_chain_id = None
    last_res_info = None  # Tuple (resSeq, iCode, resName) for TER record
    last_serial = 0

    # Check if DataFrame is empty
    if df.empty:
        buffer.write("END\n")
        content = buffer.getvalue()
        buffer.close()
        if output is not None:
            if isinstance(output, str):
                with open(output, "w") as f:
                    f.write(content)
            else:
                output.write(content)
            return None
        return content

    for _, row in df.iterrows():
        atom_data = {}

        # --- Data Extraction ---
        if format_type == "PDB":
            # Pre-process PDB values, converting None to empty strings for optional fields
            raw_alt_loc = row.get("altLoc")
            pdb_alt_loc = "" if pd.isna(raw_alt_loc) else str(raw_alt_loc)

            raw_icode = row.get("iCode")
            pdb_icode = "" if pd.isna(raw_icode) else str(raw_icode)

            raw_element = row.get("element")
            pdb_element = "" if pd.isna(raw_element) else str(raw_element)

            raw_charge = row.get("charge")
            pdb_charge = "" if pd.isna(raw_charge) else str(raw_charge)

            atom_data = {
                "record_name": row.get("record_type", "ATOM"),
                "serial": int(row.get("serial", 0)),
                "name": str(row.get("name", "")),
                "altLoc": pdb_alt_loc,
                "resName": str(row.get("resName", "")),
                "chainID": str(row.get("chainID", "")),
                "resSeq": int(row.get("resSeq", 0)),
                "iCode": pdb_icode,
                "x": float(row.get("x", 0.0)),
                "y": float(row.get("y", 0.0)),
                "z": float(row.get("z", 0.0)),
                "occupancy": float(row.get("occupancy", 1.0)),
                "tempFactor": float(row.get("tempFactor", 0.0)),
                "element": pdb_element,
                "charge": pdb_charge,
                "model": int(row.get("model", 1)),
            }
        elif format_type == "mmCIF":
            # Pre-process mmCIF values to PDB compatible format, converting None to empty strings
            raw_alt_loc = row.get("label_alt_id")
            pdb_alt_loc = "" if pd.isna(raw_alt_loc) else str(raw_alt_loc)

            raw_icode = row.get("pdbx_PDB_ins_code")
            pdb_icode = "" if pd.isna(raw_icode) else str(raw_icode)

            raw_element = row.get("type_symbol")
            pdb_element = "" if pd.isna(raw_element) else str(raw_element)

            raw_charge = row.get("pdbx_formal_charge")
            pdb_charge = "" if pd.isna(raw_charge) else str(raw_charge)

            atom_data = {
                "record_name": row.get("group_PDB", "ATOM"),
                "serial": int(row.get("id", 0)),
                "name": str(row.get("auth_atom_id", row.get("label_atom_id", ""))),
                "altLoc": pdb_alt_loc,
                "resName": str(row.get("auth_comp_id", row.get("label_comp_id", ""))),
                "chainID": str(row.get("auth_asym_id", row.get("label_asym_id"))),
                "resSeq": int(row.get("auth_seq_id", row.get("label_seq_id", 0))),
                "iCode": pdb_icode,
                "x": float(row.get("Cartn_x", 0.0)),
                "y": float(row.get("Cartn_y", 0.0)),
                "z": float(row.get("Cartn_z", 0.0)),
                "occupancy": float(row.get("occupancy", 1.0)),
                "tempFactor": float(row.get("B_iso_or_equiv", 0.0)),
                "element": pdb_element,
                "charge": pdb_charge,
                "model": int(row.get("pdbx_PDB_model_num", 1)),
            }
        else:
            raise ValueError(f"Unsupported DataFrame format: {format_type}")

        # --- MODEL/ENDMDL Records ---
        current_model_num = atom_data["model"]
        if current_model_num != last_model_num:
            if last_model_num is not None:
                buffer.write("ENDMDL\n")
            buffer.write(f"MODEL     {current_model_num:>4}\n")
            last_model_num = current_model_num
            # Reset chain/residue tracking for the new model
            last_chain_id = None
            last_res_info = None

        # --- TER Records ---
        current_chain_id = atom_data["chainID"]
        current_res_info = (
            atom_data["resSeq"],
            atom_data["iCode"],
            atom_data["resName"],
        )

        # Write TER if chain ID changes within the same model
        if last_chain_id is not None and current_chain_id != last_chain_id:
            ter_serial = str(last_serial + 1).rjust(5)
            ter_res_name = last_res_info[2].strip().rjust(3)  # Use last residue's name
            ter_chain_id = last_chain_id
            ter_res_seq = str(last_res_info[0]).rjust(4)  # Use last residue's seq num
            ter_icode = (
                last_res_info[1] if last_res_info[1] else ""
            )  # Use last residue's icode

            ter_line = f"TER   {ter_serial}      {ter_res_name} {ter_chain_id}{ter_res_seq}{ter_icode}"
            buffer.write(ter_line.ljust(80) + "\n")

        # --- Format and Write ATOM/HETATM Line ---
        pdb_line = _format_pdb_atom_line(atom_data)
        buffer.write(pdb_line + "\n")

        # --- Update Tracking Variables ---
        last_serial = atom_data["serial"]
        last_chain_id = current_chain_id
        last_res_info = current_res_info

    # --- Final Records ---
    # Add TER record for the very last chain in the last model
    if last_chain_id is not None:
        ter_serial = str(last_serial + 1).rjust(5)
        ter_res_name = last_res_info[2].strip().rjust(3)
        ter_chain_id = last_chain_id
        ter_res_seq = str(last_res_info[0]).rjust(4)
        ter_icode = last_res_info[1] if last_res_info[1] else ""

        ter_line = f"TER   {ter_serial}      {ter_res_name} {ter_chain_id}{ter_res_seq}{ter_icode}"
        buffer.write(ter_line.ljust(80) + "\n")

    # Add ENDMDL if models were used
    if last_model_num is not None:
        buffer.write("ENDMDL\n")

    buffer.write("END\n")

    # --- Output Handling ---
    content = buffer.getvalue()
    buffer.close()

    if output is not None:
        if isinstance(output, str):
            with open(output, "w") as f:
                f.write(content)
        else:
            output.write(content)
        return None
    else:
        return content