refactor: select images directly instead of container

2 files changed, 9 insertions, 10 deletions

diff --git a/src/lightbox.js b/src/lightbox.js
index 5272186..54c5283 100644
--- a/src/lightbox.js
+++ b/src/lightbox.js
@@ -11,8 +11,8 @@ export default class LightBox {
    * ============
    */
 
-  constructor({ container }) {
-    this.container = container;
+  constructor({ selector }) {
+    this.selector = selector;
     this.initOverlay();
     this.initImages();
     this.initControls();
@@ -26,7 +26,7 @@ export default class LightBox {
   }
 
   initImages() {
-    this.images = document.querySelectorAll(`${this.container} img[data-libra]`);
+    this.images = document.querySelectorAll(this.selector);
 
     function wrapInLink(image) {
       // wrap image with a link if it's not already wrapped
@@ -45,10 +45,11 @@ export default class LightBox {
     }
 
     this.images.forEach(image => {
-      // generate unique id for each image
-      image.id = crypto.randomUUID();
-      // make sure every image has a link
-      const link = wrapInLink(image);
+      if (image.tagName !== 'IMG') return;
+
+      image.id = crypto.randomUUID(); // generate unique id for each image
+      const link = wrapInLink(image); // make sure every image has a link
+
       // click link to open lightbox
       link.addEventListener('click', e => {
         e.preventDefault();
diff --git a/test/index.html b/test/index.html
index 8c1eea7..4c306b2 100644
--- a/test/index.html
+++ b/test/index.html
@@ -24,9 +24,7 @@
         </div>
         <script type="module">
             import LightBox from '/dist/libra.min.js';
-            const lightbox = new LightBox({
-              container: '.libra'
-            });
+            new LightBox({ selector: 'img[data-libra]' });
         </script>
     </body>
 </html>