Add methods to get/set metatables on Table

Also make Table::raw_get actually use raw get function instead of lua_gettable!

2 files changed, 57 insertions, 2 deletions

diff --git a/src/lua.rs b/src/lua.rs
index 13db4e7..b0a2dda 100644
--- a/src/lua.rs
+++ b/src/lua.rs
@@ -396,7 +396,7 @@ impl<'lua> Table<'lua> {
                 check_stack(lua.state, 2);
                 lua.push_ref(lua.state, &self.0);
                 lua.push_value(lua.state, key.to_lua(lua)?);
-                ffi::lua_gettable(lua.state, -2);
+                ffi::lua_rawget(lua.state, -2);
                 let res = V::from_lua(lua.pop_value(lua.state), lua)?;
                 ffi::lua_pop(lua.state, 1);
                 Ok(res)
@@ -436,6 +436,41 @@ impl<'lua> Table<'lua> {
         }
     }
 
+    pub fn get_metatable(&self) -> Option<Table<'lua>> {
+        let lua = self.0.lua;
+        unsafe {
+            stack_guard(lua.state, 0, || {
+                check_stack(lua.state, 1);
+                lua.push_ref(lua.state, &self.0);
+                if ffi::lua_getmetatable(lua.state, -1) == 0 {
+                    ffi::lua_pop(lua.state, 1);
+                    None
+                } else {
+                    let table = Table(lua.pop_ref(lua.state));
+                    ffi::lua_pop(lua.state, 1);
+                    Some(table)
+                }
+            })
+        }
+    }
+
+    pub fn set_metatable(&self, metatable: Option<Table<'lua>>) {
+        let lua = self.0.lua;
+        unsafe {
+            stack_guard(lua.state, 0, move || {
+                check_stack(lua.state, 1);
+                lua.push_ref(lua.state, &self.0);
+                if let Some(metatable) = metatable {
+                    lua.push_ref(lua.state, &metatable.0);
+                } else {
+                    ffi::lua_pushnil(lua.state);
+                }
+                ffi::lua_setmetatable(lua.state, -2);
+                ffi::lua_pop(lua.state, 1);
+            })
+        }
+    }
+
     /// Consume this table and return an iterator over the pairs of the table.
     ///
     /// This works like the Lua `pairs` function, but does not invoke the `__pairs` metamethod.
diff --git a/src/tests.rs b/src/tests.rs
index 9388df2..c3d1321 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -4,7 +4,7 @@ use std::panic::catch_unwind;
 use std::os::raw::c_void;
 
 use String as LuaString;
-use {Lua, Result, ExternalError, LightUserData, UserDataMethods, UserData, Table, Thread,
+use {Lua, Nil, Result, ExternalError, LightUserData, UserDataMethods, UserData, Table, Thread,
      ThreadStatus, Error, Function, Value, Variadic, MetaMethod};
 
 #[test]
@@ -963,6 +963,26 @@ fn test_setmetatable_gc() {
     assert_eq!(globals.get::<_, String>("val").unwrap(), "gcwascalled");
 }
 
+#[test]
+fn test_metatable() {
+    let lua = Lua::new();
+
+    let table = lua.create_table();
+    let metatable = lua.create_table();
+    metatable.set("__index", lua.create_function(|_, ()| Ok("index_value"))).unwrap();
+    table.set_metatable(Some(metatable));
+    assert_eq!(table.get::<_, String>("any_key").unwrap(), "index_value");
+    match table.raw_get::<_, Value>("any_key").unwrap() {
+        Nil => {}
+        _ => panic!(),
+    }
+    table.set_metatable(None);
+    match table.get::<_, Value>("any_key").unwrap() {
+        Nil => {}
+        _ => panic!(),
+    };
+}
+
 // Need to use compiletest-rs or similar to make sure these don't compile.
 /*
 #[test]