RNApolis 0.10.8__tar.gz → 0.11.0__tar.gz

{rnapolis-0.10.8/src/RNApolis.egg-info → rnapolis-0.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: RNApolis
-Version: 0.10.8
+Version: 0.11.0
 Summary: A Python library containing RNA-related bioinformatics functions and classes
 Home-page: https://github.com/tzok/rnapolis-py
 Author: Tomasz Zok

{rnapolis-0.10.8 → rnapolis-0.11.0}/setup.py

@@ -5,7 +5,7 @@ with open("README.md") as f:
 
 setup(
     name="RNApolis",
-    version="0.10.8",
+    version="0.11.0",
     packages=["rnapolis"],
     package_dir={"": "src"},
     author="Tomasz Zok",

{rnapolis-0.10.8 → rnapolis-0.11.0/src/RNApolis.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: RNApolis
-Version: 0.10.8
+Version: 0.11.0
 Summary: A Python library containing RNA-related bioinformatics functions and classes
 Home-page: https://github.com/tzok/rnapolis-py
 Author: Tomasz Zok

{rnapolis-0.10.8 → rnapolis-0.11.0}/src/rnapolis/molecule_filter.py

@@ -8,6 +8,7 @@ from typing import Iterable, List, Set, Tuple
 from mmcif.io.IoAdapterPy import IoAdapterPy
 from mmcif.io.PdbxReader import DataCategory, DataContainer
 
+from rnapolis.parser_v2 import parse_cif_atoms, write_pdb
 from rnapolis.util import handle_input_file
 
 # Source: https://mmcif.wwpdb.org/dictionaries/mmcif_pdbx_v50.dic/Items/_entity_poly.type.html
@@ -160,6 +161,20 @@ def filter_by_poly_types(
     return filter_cif(data, entity_ids, asym_ids, auth_asym_ids, retain_categories)
 
 
+def filter_by_entity_ids(
+    file_content: str,
+    entity_ids: Iterable[str],
+    retain_categories: Iterable[str] = ["chem_comp"],
+) -> str:
+    data = read_cif(file_content)
+    entity_ids = set(entity_ids)
+    asym_ids = select_ids(data, "struct_asym", "id", "entity_id", entity_ids)
+    auth_asym_ids = select_ids(
+        data, "atom_site", "auth_asym_id", "label_asym_id", asym_ids
+    )
+    return filter_cif(data, entity_ids, asym_ids, auth_asym_ids, retain_categories)
+
+
 def filter_by_chains(
     file_content: str,
     chains: Iterable[str],
@@ -189,6 +204,12 @@ def main():
         action="append",
         default=[],
     )
+    parser.add_argument(
+        "--filter-by-entity-ids",
+        help="filter by entity IDs, e.g. 1, 2, 3",
+        action="append",
+        default=[],
+    )
     parser.add_argument(
         "--filter-by-chains",
         help="filter by chain IDs (label_asym_id), e.g. A, B, C",
@@ -201,29 +222,43 @@ def main():
         action="append",
         default=["chem_comp"],
     )
+    parser.add_argument(
+        "--pdb", help="change output format to PDB", action="store_true"
+    )
     parser.add_argument("path", help="path to a PDBx/mmCIF file")
     args = parser.parse_args()
 
     file = handle_input_file(args.path)
+    content = None
+
     if args.filter_by_poly_types:
-        print(
-            filter_by_poly_types(
-                file.read(),
-                entity_poly_types=args.filter_by_poly_types,
-                retain_categories=args.retain_categories,
-            )
+        content = filter_by_poly_types(
+            file.read(),
+            entity_poly_types=args.filter_by_poly_types,
+            retain_categories=args.retain_categories,
+        )
+    elif args.filter_by_entity_ids:
+        content = filter_by_entity_ids(
+            file.read(),
+            entity_ids=args.filter_by_entity_ids,
+            retain_categories=args.retain_categories,
         )
     elif args.filter_by_chains:
-        print(
-            filter_by_chains(
-                file.read(),
-                chains=args.filter_by_chains,
-                retain_categories=args.retain_categories,
-            )
+        content = filter_by_chains(
+            file.read(),
+            chains=args.filter_by_chains,
+            retain_categories=args.retain_categories,
         )
     else:
         parser.print_help()
 
+    if content:
+        if args.pdb:
+            df = parse_cif_atoms(content)
+            print(write_pdb(df))
+        else:
+            print(content)
+
 
 if __name__ == "__main__":
     main()