From 63d08c2ec85c0a31acf0a90f639314f662ebcdb0 Mon Sep 17 00:00:00 2001
From: Valerij Talagayev <82884038+talagayev@users.noreply.github.com>
Date: Wed, 11 Dec 2024 01:35:56 +0100
Subject: [PATCH] Update mdonatello.py

start of addition of separate mol and pdb support
---
 mdonatello/mdonatello.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/mdonatello/mdonatello.py b/mdonatello/mdonatello.py
index 9815647..4004b37 100644
--- a/mdonatello/mdonatello.py
+++ b/mdonatello/mdonatello.py
@@ -55,7 +55,8 @@ class MoleculeVisualizer:
 
     def __init__(
         self,
-        ag: mda.core.groups.AtomGroup,
+        name: str,
+        ag: str = None,
         show_atom_indices: bool = False,
         width: int = -1,
         height: int = -1,
@@ -67,7 +68,15 @@ def __init__(
         self.height = height
         self.show_atom_indices = show_atom_indices
 
-        self.mol: Chem.Mol = ag.convert_to("RDKit")
+        self.name = name
+        self.ag = ag
+
+        if ".sdf" not in self.name and ".mol2" not in self.name:
+            self.u = mda.Universe(self.name)
+            if self.ag != None:
+                self.ag = u.select_atoms(f"resname {self.ag}")
+            self.mol: Chem.Mol = ag.convert_to("RDKit")
+            
         self.mol_noh: Chem.Mol = Chem.RemoveHs(self.mol)
         AllChem.Compute2DCoords(self.mol_noh)