fix getitem

Tumiz · May 23, 2024 · a1604f5 · a1604f5
1 parent c9e998c
commit a1604f5
Show file tree

Hide file tree

Showing 6 changed files with 635 additions and 43 deletions.
diff --git a/docs/Transform.html b/docs/Transform.html
diff --git a/docs/Vector.html b/docs/Vector.html
diff --git a/docs/Vector3.html b/docs/Vector3.html
diff --git a/notebooks/Transform.ipynb b/notebooks/Transform.ipynb
@@ -349,7 +349,7 @@
     "    [2, 7, 9, 1, 4, 8],\n",
     "    [10, 3, 8, 6, 9, 5]\n",
     "], columns=[\"x\", \"y\", \"z\", \"rvx\", \"rvy\", \"rvz\"])\n",
-    "t = py3d.Transform.from_rotation_vector(poses[[\"rvx\",\"rvy\",\"rvz\"]])@py3d.Transform.from_translation(poses[[\"x\",\"y\",\"z\"]])\n",
+    "t = py3d.Transform.from_rotation_vector(poses[\"rvx\",\"rvy\",\"rvz\"])@py3d.Transform.from_translation(poses[\"x\",\"y\",\"z\"])\n",
     "c @ t.lerp(numpy.linspace(0, 4, 20), [0, 1, 4])"
    ]
   },

diff --git a/notebooks/Vector.ipynb b/notebooks/Vector.ipynb
@@ -362,7 +362,7 @@
    "source": [
     "import py3d\n",
     "a=py3d.read_csv(\"with_cols.csv\", header=1)\n",
-    "a[\"t\",\"x\"]"
+    "a[\"t\"], a[\"t\",\"x\"], a[1,2]"
    ]
   },
   {
@@ -390,8 +390,7 @@
    "outputs": [],
    "source": [
     "import py3d\n",
-    "\n",
-    "pcd=py3d.read_pcd(\"ascii.pcd\").xyz.as_point(py3d.Color(r=1,g=0.5,a=0.3))"
+    "py3d.read_pcd(\"ascii.pcd\").xyz.as_point(py3d.Color(r=1,g=0.5,a=0.3))"
    ]
   },
   {

diff --git a/py3d/core.py b/py3d/core.py
@@ -295,14 +295,12 @@ def __imatmul__(self, value) -> Vector:
         return self @ value
 
     def __getitem__(self, keys) -> Vector:
-        if hasattr(self, "columns") and type(keys[0]) is str:
-            if type(keys) is str:
-                i = self.columns.index(keys)
-            else:
-                i = [self.columns.index(key) for key in keys]
-            return self[..., i]
-        else:
-            return super().__getitem__(keys)
+        if hasattr(self, "columns"):
+            if isinstance(keys, str):
+                keys = ..., self.columns.index(keys)
+            elif isinstance(keys, tuple) and all(isinstance(k, str) for k in keys):
+                keys = ..., [self.columns.index(key) for key in keys]
+        return super().__getitem__(keys)
 
     def tile(self, *n) -> Vector:
         return numpy.tile(self, n + self.ndim * (1,))