File size: 3,757 Bytes
9f27665
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
"""
Example Data Loader for RNA Motif Comparison Tool
Provides functionality to load example PDB files from data folder
"""

import os
from pathlib import Path
import streamlit as st

def get_example_pdbs(data_folder="data"):
    """
    Get list of example PDB files from the data folder.
    
    Args:
        data_folder: Path to folder containing example PDB files
        
    Returns:
        Dictionary with filename as key and full path as value
    """
    examples = {}
    
    if not os.path.exists(data_folder):
        return examples
    
    # Get all PDB files in the data folder
    data_path = Path(data_folder)
    for pdb_file in data_path.glob("*.pdb"):
        examples[pdb_file.name] = str(pdb_file)
    
    # Also check for uppercase .PDB extension
    for pdb_file in data_path.glob("*.PDB"):
        examples[pdb_file.name] = str(pdb_file)
    
    return examples


def create_example_selector(data_folder="data"):
    """
    Create a Streamlit interface for selecting example PDB files.
    
    Args:
        data_folder: Path to folder containing example PDB files
        
    Returns:
        List of selected file paths
    """
    examples = get_example_pdbs(data_folder)
    
    if not examples:
        st.warning(f"⚠️ No example PDB files found in '{data_folder}/' folder")
        return []
    
    st.info(f"📁 Found {len(examples)} example PDB files in '{data_folder}/' folder")
    
    # Create multiselect for choosing examples
    selected_names = st.multiselect(
        "Select example PDB files to load",
        options=sorted(examples.keys()),
        help="Choose one or more example structures"
    )
    
    # Return full paths of selected files
    selected_paths = [examples[name] for name in selected_names]
    
    if selected_paths:
        st.success(f"✅ Selected {len(selected_paths)} example file(s)")
    
    return selected_paths


def load_example_as_uploaded_file(file_path):
    """
    Load a PDB file and convert it to a format similar to Streamlit's UploadedFile.
    
    Args:
        file_path: Path to the PDB file
        
    Returns:
        File-like object with name and getbuffer() method
    """
    class MockUploadedFile:
        def __init__(self, path):
            self.name = os.path.basename(path)
            self.path = path
            with open(path, 'rb') as f:
                self._content = f.read()
        
        def getbuffer(self):
            return self._content
        
        def read(self):
            return self._content
    
    return MockUploadedFile(file_path)


def get_example_info(data_folder="data"):
    """
    Get information about example PDB files.
    
    Args:
        data_folder: Path to folder containing example PDB files
        
    Returns:
        Dictionary with file info
    """
    examples = get_example_pdbs(data_folder)
    info = {}
    
    for name, path in examples.items():
        try:
            with open(path, 'r') as f:
                lines = f.readlines()
            
            # Count atoms/residues
            atom_count = sum(1 for line in lines if line.startswith('ATOM') or line.startswith('HETATM'))
            
            # Get header info if available
            header = ""
            for line in lines:
                if line.startswith('HEADER'):
                    header = line[10:].strip()
                    break
            
            info[name] = {
                'path': path,
                'atoms': atom_count,
                'header': header,
                'lines': len(lines)
            }
        except Exception as e:
            info[name] = {
                'path': path,
                'error': str(e)
            }
    
    return info