Return undef if column number is exceeded in add_trait()

bioperl · May 21, 2012 · bafabb5 · bafabb5
1 parent 791f66a
commit bafabb5
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 24 deletions.
diff --git a/Bio/Tree/TreeFunctionsI.pm b/Bio/Tree/TreeFunctionsI.pm
@@ -1134,10 +1134,11 @@ sub move_id_to_bootstrap{
            method. This means that you can retrieve the trait values using the
            get_tag_values() method (see the documentation for Bio::Tree::Node).
            If a leaf node has no value in the trait file, an exception is thrown.
- Returns : Trait name (a scalar)
+ Returns : Trait name (a scalar) on success, undef on failure (for example, if
+           the column index requested was too large).
  Args    : Name of trait file (scalar string)
            Index of trait file column (scalar int). Note that numbering starts
-               at 0. Default: 1 (second column)
+           at 0. Default: 1 (second column)
 
 =cut
 
@@ -1146,46 +1147,47 @@ sub _read_trait_file {
     my $file = shift;
     my $column = shift || 1;
 
+    my $trait_name;
     my $traits;
     open my $TRAIT, "<", $file or $self->("Could not open file $file: $!\n");
 
     my $first_line = 1;
     while (<$TRAIT>) {
+        s/['"]//g;
+        my @line = split;
         if ($first_line) {
             $first_line = 0;
-            s/['"]//g;
-            my @line = split;
-            $traits->{'my_trait_name'} = $line[$column];
+            $trait_name = $line[$column];
             next;
         }
-        s/['"]//g;
-        my @line = split;
         last unless $line[0];
         $traits->{$line[0]} = $line[$column];
     }
 
     close $TRAIT;
-    return $traits;
+    return $trait_name, $traits;
 }
 
 sub add_trait {
     my $self = shift;
     my $file = shift;
     my $column = shift;
 
-    my $traits = $self->_read_trait_file($file, $column); # filename, trait column
-    my $key = $traits->{'my_trait_name'};
-    foreach my $node ($self->get_leaf_nodes) {
-        # strip quotes from the node id
-        $node->id($1) if $node->id =~ /^['"]+(.*)['"]+$/;
-        eval {
-            $node->verbose(2);
-            $node->add_tag_value($key, $traits->{ $node->id } );
-        };
-        $self->throw("No trait for node [".$node->id."/".$node->internal_id."]")
-            if $@;
+    my ($trait_name, $traits) = $self->_read_trait_file($file, $column);
+
+    if (defined $trait_name) {
+        for my $node ($self->get_leaf_nodes) {
+            # strip quotes from the node id
+            $node->id($1) if $node->id =~ /^['"]+(.*)['"]+$/;
+            eval {
+                $node->verbose(2);
+                $node->add_tag_value($trait_name, $traits->{ $node->id } );
+            };
+            $self->throw("No trait for node [".$node->id."/".$node->internal_id."]")
+               if $@;
+        }
     }
-    return $key;
+    return $trait_name;
 }
 
 

diff --git a/t/Tree/TreeStatistics.t b/t/Tree/TreeStatistics.t
@@ -7,7 +7,7 @@ BEGIN {
     use lib '.';
     use Bio::Root::Test;
 
-    test_begin(-tests => 40);
+    test_begin(-tests => 41);
 
     use_ok('Bio::TreeIO');
     use_ok('Bio::Tree::Statistics');
@@ -16,7 +16,7 @@ BEGIN {
 
 use Data::Dumper;
 my $in = Bio::TreeIO->new(-format => 'nexus',
-			  -file   => test_input_file('traittree.nexus'));
+                          -file   => test_input_file('traittree.nexus'));
 my $tree = $in->next_tree;
 my $node = $tree->find_node(-id => 'N14');
 
@@ -26,8 +26,11 @@ is $stats->cherries($tree), 8, 'cherries';
 is $stats->cherries($tree, $node), 4, 'cherries';
 
 # traits
-my $key = $tree->add_trait(test_input_file('traits.tab'), 3);
-is ($key, 'intermediate', 'read traits');
+my $key = $tree->add_trait(test_input_file('traits.tab'), 4);
+is $key, undef, 'read traits'; # exceeded column number
+
+$key = $tree->add_trait(test_input_file('traits.tab'), 3);
+is $key, 'intermediate';
 
 is $stats->ps($tree, $key), 4, 'parsimony score';
 is $stats->ps($tree, $key, $node), 1, 'subtree parsimony score';