about summary refs log tree commit
diff options
context:
space:
mode:
authorEric Wong <normalperson@yhbt.net>2010-09-02 16:34:45 -0700
committerEric Wong <normalperson@yhbt.net>2010-09-02 16:34:45 -0700
commit84be9dc370fe32f05cbec7ff7da7f774e4ba12af (patch)
treea5eb4703deecc84d562072f7c9b1760591c5ebd3
parent33eb621b54552e5398752976f07ef9494a610841 (diff)
downloadrpatricia-84be9dc370fe32f05cbec7ff7da7f774e4ba12af.tar.gz
show_nodes may be given a custom IO-like object
This makes it easy to dump the node output to an
alternate IO/StringIO object without temporarily
replacing $stdout.
-rw-r--r--ext/rpatricia/rpatricia.c11
-rw-r--r--test/test_show_nodes.rb22
2 files changed, 30 insertions, 3 deletions
diff --git a/ext/rpatricia/rpatricia.c b/ext/rpatricia/rpatricia.c
index 38246e1..1c81ba1 100644
--- a/ext/rpatricia/rpatricia.c
+++ b/ext/rpatricia/rpatricia.c
@@ -170,7 +170,7 @@ rb_str_set_len(VALUE str, long len)
 #endif
 
 static VALUE
-p_print_nodes (VALUE self)
+p_print_nodes (int argc, VALUE *argv, VALUE self)
 {
   ID id_printf = rb_intern("printf");
   VALUE fmt = rb_str_new2("node: %s\n");
@@ -178,15 +178,20 @@ p_print_nodes (VALUE self)
   char *cbuf;
   patricia_tree_t *tree;
   patricia_node_t *node;
+  VALUE out;
   Data_Get_Struct(self, patricia_tree_t, tree);
 
+  rb_scan_args(argc, argv, "01", &out);
+  if (NIL_P(out))
+    out = rb_stdout;
+
   if (tree->head) {
     PATRICIA_WALK(tree->head, node) {
       rb_str_resize(buf, 128);
       cbuf = RSTRING_PTR(buf);
       prefix_toa2x(node->prefix, cbuf, 1);
       rb_str_set_len(buf, strlen(cbuf));
-      rb_funcall(rb_stdout, id_printf, 2, fmt, buf);
+      rb_funcall(out, id_printf, 2, fmt, buf);
     } PATRICIA_WALK_END;
   }
   return Qtrue;
@@ -328,7 +333,7 @@ Init_rpatricia (void)
 
   /* derivatives of climb */
   rb_define_method(cPatricia, "num_nodes", p_num_nodes, 0);
-  rb_define_method(cPatricia, "show_nodes", p_print_nodes, 0);
+  rb_define_method(cPatricia, "show_nodes", p_print_nodes, -1);
 
   /* destroy tree */
   rb_define_method(cPatricia, "destroy", p_destroy, 0);
diff --git a/test/test_show_nodes.rb b/test/test_show_nodes.rb
index b4aa3a0..91e6f81 100644
--- a/test/test_show_nodes.rb
+++ b/test/test_show_nodes.rb
@@ -29,4 +29,26 @@ class TestShowNodes < Test::Unit::TestCase
     ].join("\n") << "\n"
     assert_equal expect, tmpout.string
   end
+
+  def test_show_nodes_specify_io
+    t = Patricia.new
+    assert_nothing_raised do
+      t.add("127.0.0.0/24")
+      t.add("192.168.1.0/24")
+      t.add("192.168.2.0/24")
+      t.add("192.168.3.100")
+      t.add("10.0.0.0/8", "pref_10")
+    end
+    tmpout = StringIO.new
+    assert_nothing_raised { t.show_nodes(tmpout) }
+    expect = [
+      "node: 10.0.0.0/8",
+      "node: 127.0.0.0/24",
+      "node: 192.168.1.0/24",
+      "node: 192.168.2.0/24",
+      "node: 192.168.3.100/32"
+    ].join("\n") << "\n"
+    assert_equal expect, tmpout.string
+  end
+
 end